Update app.py

app.py

@@ -1,41 +1,64 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
-import torch
+import cv2
+import numpy as np
+import re
 
 # Lade das Modell
 model_path = hf_hub_download(repo_id="foduucom/stockmarket-pattern-detection-yolov8", filename="model.pt")
 model = YOLO(model_path)
 
 def analyze_image(image, prompt):
-    # Verwende den Prompt (falls nötig, hier als Kontext für die Verarbeitung)
-    # YOLOv8 ignoriert den Prompt direkt, daher speichern wir ihn für die Logik
-    results = model.predict(source=image, save=False)
+    # Konvertiere PIL-Bild zu OpenCV-Format
+    image_np = np.array(image)
+    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+
+    # Führe Objekterkennung durch
+    results = model.predict(source=image_np, save=False)
+    
+    # Extrahiere Kerzen
     detections = []
     for result in results:
         for box in result.boxes:
             label = result.names[int(box.cls)]
             confidence = float(box.conf)
-            # Farben basierend auf Label oder Prompt (z. B. "bullish" für grün)
-            color = "green" if "bullish" in prompt.lower() or "Bullish" in label else "red"
+            # Extrahiere Bounding-Box-Koordinaten
+            xmin, ymin, xmax, ymax = box.xyxy[0].tolist()
+            
+            # Schneide die Kerze aus dem Bild für Farbanalyse
+            candle_roi = image_cv[int(ymin):int(ymax), int(xmin):int(xmax)]
+            # Berechne dominante Farbe (RGB)
+            mean_color = np.mean(candle_roi, axis=(0, 1)).astype(int)
+            color_rgb = f"RGB({mean_color[2]},{mean_color[1]},{mean_color[0]})"  # RGB-Format
+            
             detections.append({
                 "pattern": label,
                 "confidence": confidence,
-                "color": color,
-                "prompt_used": prompt  # Rückgabe des Prompts zur Überprüfung
+                "color": color_rgb,
+                "x_center": (xmin + xmax) / 2,  # Für Sortierung nach Position
+                "prompt_used": prompt
             })
+
+    # Sortiere Kerzen nach x-Position (von rechts nach links = neueste zuerst)
+    detections = sorted(detections, key=lambda x: x["x_center"], reverse=True)
+    
+    # Begrenze auf die letzten 10 Kerzen, wenn im Prompt gefordert
+    if "last 10 candles" in prompt.lower():
+        detections = detections[:10]
+    
     return detections
 
-# Erstelle Gradio-Schnittstelle mit Bild- und Text-Eingabe
+# Erstelle Gradio-Schnittstelle
 iface = gr.Interface(
     fn=analyze_image,
     inputs=[
         gr.Image(type="pil", label="Upload TradingView Screenshot"),
-        gr.Textbox(label="Prompt", placeholder="Enter your prompt, e.g., 'Detect candlestick patterns and colors'")
+        gr.Textbox(label="Prompt", placeholder="Enter your prompt, e.g., 'List last 10 candles with their colors'")
     ],
     outputs="json",
     title="Candlestick Pattern Detection",
-    description="Upload a TradingView screenshot and provide a prompt to detect candlestick patterns and colors."
+    description="Upload a TradingView screenshot and provide a prompt to detect candlestick patterns and their colors."
 )
 
 # Starte die App