Update app.py

app.py

@@ -1,82 +1,101 @@
 import gradio as gr
-from transformers import AutoModel, AutoImageProcessor
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
 import cv2
 import numpy as np
-import torch
-import os
-from huggingface_hub import login
+import easyocr
+import logging
 
-# Authentifizierung mit Hugging Face API-Token
-HF_TOKEN = os.getenv("HF_TOKEN")
-if HF_TOKEN:
-    login(HF_TOKEN)
-else:
-    raise ValueError("HF_TOKEN environment variable not set. Please add it in Space settings.")
+# Logging einrichten
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Lade das Modell und den Image Processor
-model = AutoModel.from_pretrained("facebook/dinov3-convnext-small-pretrain-lvd1689m")
-image_processor = AutoImageProcessor.from_pretrained("facebook/dinov3-convnext-small-pretrain-lvd1689m")
+# Lade das Modell
+model_path = hf_hub_download(repo_id="foduucom/stockmarket-pattern-detection-yolov8", filename="model.pt")
+model = YOLO(model_path)
+
+# OCR für Preise
+reader = easyocr.Reader(['en'], gpu=False)
 
 def analyze_image(image, prompt):
+    logger.info("Starting image analysis with prompt: %s", prompt)
+    
     # Konvertiere PIL-Bild zu OpenCV-Format
     image_np = np.array(image)
     image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+    
+    # Bildvorverarbeitung: Kontrast erhöhen
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    gray = cv2.cvtColor(image_cv, cv2.COLOR_BGR2GRAY)
+    enhanced = clahe.apply(gray)
+    image_cv = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR)
+    logger.info("Image preprocessed: shape=%s", image_np.shape)
 
-    # Extrahiere Features mit DINOv3
-    inputs = image_processor(images=image_np, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
+    # Führe Objekterkennung durch
+    results = model.predict(source=image_np, conf=0.3, iou=0.5, save=False)
+    logger.info("YOLO predictions: %d boxes detected", len(results[0].boxes))
 
-    # Verbesserte Bildanalyse mit OpenCV
-    gray = cv2.cvtColor(image_cv, cv2.COLOR_BGR2GRAY)
-    # Adaptiver Schwellenwert für bessere Konturenerkennung
-    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
-    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    # Extrahiere Kerzen
+    detections = []
+    for result in results:
+        for box in result.boxes:
+            label = result.names[int(box.cls)]
+            confidence = float(box.conf)
+            xmin, ymin, xmax, ymax = box.xyxy[0].tolist()
+            logger.info("Detected: %s, confidence=%.2f, box=(%.0f, %.0f, %.0f, %.0f)", 
+                       label, confidence, xmin, ymin, xmax, ymax)
+
+            # Extrahiere Farbe (Fokus auf Kerzenkörper)
+            candle_roi = image_cv[int(ymin):int(ymax), int(xmin):int(xmax)]
+            if candle_roi.size == 0:
+                logger.warning("Empty ROI for box: (%.0f, %.0f, %.0f, %.0f)", xmin, ymin, xmax, ymax)
+                continue
+            mean_color = np.mean(candle_roi, axis=(0, 1)).astype(int)
+            color_rgb = f"RGB({mean_color[2]},{mean_color[1]},{mean_color[0]})"
+
+            # OCR für Preise (erweitere ROI für Achsen)
+            price_roi = image_cv[max(0, int(ymin)-50):min(image_np.shape[0], int(ymax)+50), 
+                                max(0, int(xmin)-50):min(image_np.shape[1], int(xmax)+50)]
+            price_text = reader.readtext(price_roi, detail=0, allowlist='0123456789.')
+            prices = ' '.join(price_text) if price_text else "No price detected"
+            logger.info("OCR prices: %s", prices)
+
+            detections.append({
+                "pattern": label,
+                "confidence": confidence,
+                "color": color_rgb,
+                "prices": prices,
+                "x_center": (xmin + xmax) / 2
+            })
 
-    # Analysiere das Bild basierend auf dem Prompt
-    description = []
-    if "what do you see" in prompt.lower() or "was siehst du" in prompt.lower():
-        if len(contours) == 0:
-            description.append("Das Bild enthält keine klar erkennbaren Objekte.")
-        else:
-            for idx, contour in enumerate(contours[:20]):  # Begrenze auf 20 Objekte
-                if cv2.contourArea(contour) < 200 or cv2.contourArea(contour) > (image_np.shape[0] * image_np.shape[1] * 0.5):  # Filtere kleine/große Konturen
-                    continue
-                x, y, w, h = cv2.boundingRect(contour)
-                # Extrahiere Farbe der Region
-                roi = image_cv[y:y+h, x:x+w]
-                if roi.size == 0:
-                    continue
-                mean_color = np.mean(roi, axis=(0, 1)).astype(int)
-                color_rgb = f"RGB({mean_color[2]},{mean_color[1]},{mean_color[0]})"
-                # Größenkategorie
-                size = "small" if w * h < 1000 else "medium" if w * h < 5000 else "large"
-                description.append({
-                    "object": f"Object_{idx}",
-                    "color": color_rgb,
-                    "position": f"x={x}, y={y}, width={w}, height={h}",
-                    "size": size
-                })
+    # Sortiere nach x-Position (rechts nach links = neueste Kerzen)
+    detections = sorted(detections, key=lambda x: x["x_center"], reverse=True)
+    logger.info("Sorted detections: %d", len(detections))
 
-    # Einfache Analyse der DINOv3-Features (z. B. Anzahl der Feature-Regionen)
-    feature_info = str(outputs.last_hidden_state.shape) if hasattr(outputs, 'last_hidden_state') else "No features extracted."
+    # Begrenze auf die letzten 8 Kerzen
+    if "last 8 candles" in prompt.lower() or "letzte 8 kerzen" in prompt.lower():
+        detections = detections[:8]
 
+    # Debugging: Wenn leer, gib Hinweis
+    if not detections:
+        logger.warning("No detections found. Check image quality or model configuration.")
+        return {"prompt": prompt, "description": "No candlesticks detected. Ensure clear image and visible candles."}
+    
     return {
         "prompt": prompt,
-        "description": description if description else "No objects detected.",
-        "features_shape": feature_info
+        "detections": detections
     }
 
 # Erstelle Gradio-Schnittstelle
 iface = gr.Interface(
     fn=analyze_image,
     inputs=[
-        gr.Image(type="pil", label="Upload an Image"),
-        gr.Textbox(label="Prompt", placeholder="Enter your prompt, e.g., 'Was siehst du auf dem Bild?'")
+        gr.Image(type="pil", label="Upload TradingView Screenshot"),
+        gr.Textbox(label="Prompt", placeholder="Enter your prompt, e.g., 'List last 8 candles with their colors'")
     ],
     outputs="json",
-    title="General Image Analysis with DINOv3",
-    description="Upload an image and provide a prompt to get a description of what the model sees."
+    title="Stock Chart Analysis with YOLOv8",
+    description="Upload a TradingView screenshot to detect the last 8 candlesticks, their colors, and prices."
 )
 
 iface.launch()