import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import torch

# 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):
    # Führe Objekterkennung durch
    results = model.predict(source=image, save=False)
    # Extrahiere Ergebnisse
    detections = []
    for result in results:
        for box in result.boxes:
            label = result.names[int(box.cls)]
            confidence = float(box.conf)
            detections.append({
                "pattern": label,
                "confidence": confidence,
                "color": "green" if "Bullish" in label else "red"
            })
    return detections

# Erstelle Gradio-Schnittstelle
iface = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="Candlestick Pattern Detection",
    description="Upload a TradingView screenshot to detect candlestick patterns and colors."
)

# Starte die App
iface.launch()