Update app.py

app.py

@@ -2,52 +2,52 @@ import os, time, json
 import numpy as np
 import gradio as gr
 import torch
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 import cv2
 from transformers import AutoImageProcessor, RTDetrForObjectDetection
+from ultralytics import YOLO
+from huggingface_hub import hf_hub_download
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Globale Modelle
 rtdetr_r50_model = None
 rtdetr_r50_processor = None
-rtdetr_r101_model = None
-rtdetr_r101_processor = None
+yolo_o365_model = None
 
 def load_models():
-    global rtdetr_r50_model, rtdetr_r50_processor, rtdetr_r101_model, rtdetr_r101_processor
+    global rtdetr_r50_model, rtdetr_r50_processor, yolo_o365_model
     
-    print("Lade RT-DETR R50 (COCO + Objects365)...")
+    print("Lade RT-DETR R50 (COCO 80 Klassen)...")
     model_id_r50 = "PekingU/rtdetr_r50vd_coco_o365"
     try:
         rtdetr_r50_model = RTDetrForObjectDetection.from_pretrained(model_id_r50).to(DEVICE)
         rtdetr_r50_processor = AutoImageProcessor.from_pretrained(model_id_r50)
-        print(f"R50 geladen - kennt {len(rtdetr_r50_model.config.id2label)} Klassen")
-        
-        # Debug: Zeige einige Labels
-        labels = list(rtdetr_r50_model.config.id2label.values())[:20]
-        print(f"Erste 20 Labels: {labels}")
-        
-        # Suche nach Badezimmer-relevanten Labels
-        bad_labels = [l for l in rtdetr_r50_model.config.id2label.values() 
-                      if any(word in l.lower() for word in ['toilet', 'sink', 'faucet', 'mirror', 'towel', 'bath'])]
-        print(f"Badezimmer-relevante Labels gefunden: {bad_labels}")
-        
+        print(f"RT-DETR R50 geladen - kennt {len(rtdetr_r50_model.config.id2label)} Klassen")
     except Exception as e:
         return f"Fehler beim Laden von RT-DETR R50: {str(e)}"
     
-    print("Lade RT-DETR R101 (COCO + Objects365)...")
-    model_id_r101 = "PekingU/rtdetr_r101vd_coco_o365"
+    print("Lade YOLO11n (Objects365 - 365 Klassen)...")
     try:
-        rtdetr_r101_model = RTDetrForObjectDetection.from_pretrained(model_id_r101).to(DEVICE)
-        rtdetr_r101_processor = AutoImageProcessor.from_pretrained(model_id_r101)
-        print(f"R101 geladen - kennt {len(rtdetr_r101_model.config.id2label)} Klassen")
+        weights_path = hf_hub_download("NRtred/yolo11n_object365", "yolo11n_object365.pt")
+        yolo_o365_model = YOLO(weights_path)
+        print(f"YOLO11n geladen - kennt {len(yolo_o365_model.names)} Klassen")
+        
+        # Debug: Zeige einige YOLO Labels
+        yolo_labels = list(yolo_o365_model.names.values())[:30]
+        print(f"Erste 30 YOLO Labels: {yolo_labels}")
+        
+        # Suche nach Badezimmer-relevanten Labels in YOLO
+        bad_labels_yolo = [l for l in yolo_o365_model.names.values() 
+                           if any(word in l.lower() for word in ['toilet', 'sink', 'faucet', 'mirror', 'towel', 'bath', 'shower'])]
+        print(f"YOLO Badezimmer-Labels: {bad_labels_yolo}")
+        
     except Exception as e:
-        return f"Fehler beim Laden von RT-DETR R101: {str(e)}"
+        return f"Fehler beim Laden von YOLO11n: {str(e)}"
     
-    return f"Beide Modelle geladen! R50 und R101 mit COCO+Objects365 Training"
+    return f"Beide Modelle geladen! RT-DETR R50 (80 COCO) und YOLO11n (365 Objects365)"
 
-def detect_with_rtdetr(image: Image.Image, model, processor, model_name: str, confidence_threshold=0.25):
+def detect_with_rtdetr(image: Image.Image, model, processor, confidence_threshold=0.25):
     start = time.time()
     
     inputs = processor(images=image, return_tensors="pt").to(DEVICE)
@@ -68,8 +68,7 @@ def detect_with_rtdetr(image: Image.Image, model, processor, model_name: str, co
     
     id2label = model.config.id2label if hasattr(model.config, 'id2label') else {}
     
-    # Farbe je nach Modell
-    box_color = "red" if "r50" in model_name.lower() else "blue"
+    box_color = "red"
     
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
         x1, y1, x2, y2 = [float(x) for x in box.tolist()]
@@ -89,67 +88,105 @@ def detect_with_rtdetr(image: Image.Image, model, processor, model_name: str, co
     dur = time.time() - start
     return annotated, detections, dur
 
+def detect_with_yolo(image: Image.Image, model, confidence_threshold=0.25):
+    start = time.time()
+    
+    # YOLO inference
+    results = model(image, conf=confidence_threshold, device=DEVICE.type)
+    
+    detections = []
+    annotated = image.copy()
+    draw = ImageDraw.Draw(annotated)
+    
+    box_color = "blue"
+    
+    for r in results:
+        boxes = r.boxes
+        if boxes is not None:
+            for box in boxes:
+                x1, y1, x2, y2 = box.xyxy[0].tolist()
+                conf = float(box.conf[0])
+                cls = int(box.cls[0])
+                label = model.names[cls]
+                
+                detections.append({
+                    "label": label,
+                    "confidence": round(conf, 3),
+                    "bbox": [int(x1), int(y1), int(x2), int(y2)]
+                })
+                
+                draw.rectangle([x1, y1, x2, y2], outline=box_color, width=3)
+                draw.text((x1, max(0, y1 - 14)), f"{label}: {conf:.2f}", fill=box_color)
+    
+    dur = time.time() - start
+    return annotated, detections, dur
+
 def compare_models(image: Image.Image, confidence_threshold: float):
     if image is None:
         return None, None, "Bitte lade ein Bild hoch."
 
-    # R50 Detection (schneller, weniger genau)
-    r50_img, r50_det, r50_t = detect_with_rtdetr(
-        image, rtdetr_r50_model, rtdetr_r50_processor, "R50", confidence_threshold
+    # RT-DETR R50 Detection (80 COCO Klassen)
+    rtdetr_img, rtdetr_det, rtdetr_t = detect_with_rtdetr(
+        image, rtdetr_r50_model, rtdetr_r50_processor, confidence_threshold
     )
     
-    # R101 Detection (langsamer, genauer)
-    r101_img, r101_det, r101_t = detect_with_rtdetr(
-        image, rtdetr_r101_model, rtdetr_r101_processor, "R101", confidence_threshold
+    # YOLO11n Detection (365 Objects365 Klassen)
+    yolo_img, yolo_det, yolo_t = detect_with_yolo(
+        image, yolo_o365_model, confidence_threshold
     )
 
     # Labels extrahieren und zählen
-    r50_objects = {}
-    for d in r50_det:
+    rtdetr_objects = {}
+    for d in rtdetr_det:
         label = d["label"]
-        if label not in r50_objects:
-            r50_objects[label] = 0
-        r50_objects[label] += 1
+        if label not in rtdetr_objects:
+            rtdetr_objects[label] = 0
+        rtdetr_objects[label] += 1
     
-    r101_objects = {}
-    for d in r101_det:
+    yolo_objects = {}
+    for d in yolo_det:
         label = d["label"]
-        if label not in r101_objects:
-            r101_objects[label] = 0
-        r101_objects[label] += 1
+        if label not in yolo_objects:
+            yolo_objects[label] = 0
+        yolo_objects[label] += 1
     
     # Alle einzigartigen Labels
-    all_labels = set(list(r50_objects.keys()) + list(r101_objects.keys()))
+    all_labels = set(list(rtdetr_objects.keys()) + list(yolo_objects.keys()))
     
     # Nur in einem Modell gefunden
-    nur_r50 = set(r50_objects.keys()) - set(r101_objects.keys())
-    nur_r101 = set(r101_objects.keys()) - set(r50_objects.keys())
-    beide = set(r50_objects.keys()) & set(r101_objects.keys())
+    nur_rtdetr = set(rtdetr_objects.keys()) - set(yolo_objects.keys())
+    nur_yolo = set(yolo_objects.keys()) - set(rtdetr_objects.keys())
+    beide = set(rtdetr_objects.keys()) & set(yolo_objects.keys())
     
     # Vergleichstabelle
-    comparison_table = "| Objekt | R50 (schnell) | R101 (genau) | Differenz |\n"
-    comparison_table += "|--------|---------------|--------------|----------|\n"
+    comparison_table = "| Objekt | RT-DETR (80) | YOLO (365) | Anmerkung |\n"
+    comparison_table += "|--------|--------------|------------|----------|\n"
     
     for label in sorted(all_labels):
-        r50_count = r50_objects.get(label, 0)
-        r101_count = r101_objects.get(label, 0)
-        diff = r50_count - r101_count
-        diff_str = f"+{diff}" if diff > 0 else str(diff)
-        comparison_table += f"| {label} | {r50_count} | {r101_count} | {diff_str} |\n"
+        rtdetr_count = rtdetr_objects.get(label, 0)
+        yolo_count = yolo_objects.get(label, 0)
+        note = ""
+        if label in nur_rtdetr:
+            note = "⚠️ Nur COCO"
+        elif label in nur_yolo:
+            note = "✨ O365 Extra"
+        comparison_table += f"| {label} | {rtdetr_count} | {yolo_count} | {note} |\n"
     
     # Vollständige JSON Ausgabe
     full_json = json.dumps({
         "rtdetr_r50": {
-            "total_objects": len(r50_det),
-            "inference_time_ms": round(r50_t * 1000, 1),
-            "unique_finds": list(nur_r50),
-            "all_detections": r50_det
+            "total_objects": len(rtdetr_det),
+            "inference_time_ms": round(rtdetr_t * 1000, 1),
+            "unique_finds": list(nur_rtdetr),
+            "available_classes": 80,
+            "all_detections": rtdetr_det
         },
-        "rtdetr_r101": {
-            "total_objects": len(r101_det),
-            "inference_time_ms": round(r101_t * 1000, 1),
-            "unique_finds": list(nur_r101),
-            "all_detections": r101_det
+        "yolo11n_o365": {
+            "total_objects": len(yolo_det),
+            "inference_time_ms": round(yolo_t * 1000, 1),
+            "unique_finds": list(nur_yolo),
+            "available_classes": 365,
+            "all_detections": yolo_det
         },
         "gemeinsam": {
             "beide_gefunden": list(beide),
@@ -157,24 +194,24 @@ def compare_models(image: Image.Image, confidence_threshold: float):
         }
     }, ensure_ascii=False, indent=2)
 
-    # Markdown-String ohne problematische Backticks in f-string
-    md = "## RT-DETR Modellvergleich: R50 vs R101 (beide COCO + Objects365)\n\n"
-    md += "### Training-Basis\n"
-    md += "Beide Modelle wurden auf **COCO (80 Klassen) + Objects365 (365 Klassen)** trainiert.\n"
-    md += "Das bedeutet sie kennen theoretisch **400+ verschiedene Objekttypen**!\n\n"
+    # Markdown-String
+    md = "## Modellvergleich: RT-DETR R50 (COCO) vs YOLO11n (Objects365)\n\n"
+    md += "### Klassen-Unterschied\n"
+    md += "- **RT-DETR R50:** 80 COCO Klassen (Standard-Objekte)\n"
+    md += "- **YOLO11n:** 365 Objects365 Klassen (detaillierte Objekterkennung)\n\n"
     md += "### Zusammenfassung\n"
-    md += f"- **RT-DETR R50:** {len(r50_det)} Objekte in {r50_t*1000:.1f}ms\n"
-    md += f"- **RT-DETR R101:** {len(r101_det)} Objekte in {r101_t*1000:.1f}ms\n"
-    md += f"- **Geschwindigkeitsfaktor:** R50 ist {r101_t/r50_t:.2f}x schneller\n\n"
+    md += f"- **RT-DETR R50:** {len(rtdetr_det)} Objekte in {rtdetr_t*1000:.1f}ms\n"
+    md += f"- **YOLO11n O365:** {len(yolo_det)} Objekte in {yolo_t*1000:.1f}ms\n"
+    md += f"- **Geschwindigkeitsfaktor:** {'RT-DETR' if rtdetr_t < yolo_t else 'YOLO'} ist {max(rtdetr_t, yolo_t)/min(rtdetr_t, yolo_t):.2f}x schneller\n\n"
     md += "### Exklusive Funde\n"
-    md += f"- **Nur R50 gefunden:** {', '.join(nur_r50) if nur_r50 else 'Keine'}\n"
-    md += f"- **Nur R101 gefunden:** {', '.join(nur_r101) if nur_r101 else 'Keine'}\n"
+    md += f"- **Nur RT-DETR (COCO):** {', '.join(nur_rtdetr) if nur_rtdetr else 'Keine'}\n"
+    md += f"- **Nur YOLO (O365 Extra):** {', '.join(nur_yolo) if nur_yolo else 'Keine'}\n"
     md += f"- **Beide gefunden:** {len(beide)} gemeinsame Objekttypen\n\n"
     md += "### Detaillierter Vergleich\n"
     md += comparison_table + "\n"
-    md += "### Modell-Unterschiede\n"
-    md += "- **R50:** ResNet-50 Backbone, 50 Layer tief, schneller aber weniger genau\n"
-    md += "- **R101:** ResNet-101 Backbone, 101 Layer tief, genauer aber langsamer\n\n"
+    md += "### Modell-Eigenschaften\n"
+    md += "- **RT-DETR:** Transformer-basiert, End-to-End Detection, COCO-fokussiert\n"
+    md += "- **YOLO11n:** CNN-basiert, ultraschnell, 365 detaillierte Objektklassen\n\n"
     md += "### Alle Erkennungen (JSON)\n"
     md += "<details>\n"
     md += "<summary>Klick für vollständige Daten</summary>\n\n"
@@ -183,7 +220,7 @@ def compare_models(image: Image.Image, confidence_threshold: float):
     md += "\n```\n"
     md += "</details>\n"
 
-    return r50_img, r101_img, md
+    return rtdetr_img, yolo_img, md
 
 # Modelle beim Start laden
 print("Starte Modell-Ladevorgang...")
@@ -191,9 +228,9 @@ load_status = load_models()
 print(load_status)
 
 # Gradio Interface
-with gr.Blocks(title="RT-DETR Modellvergleich") as demo:
-    gr.Markdown("# 🔍 RT-DETR Modellvergleich: R50 vs R101")
-    gr.Markdown("Vergleiche die Performance von RT-DETR ResNet-50 und ResNet-101 Modellen (beide trainiert auf COCO + Objects365)")
+with gr.Blocks(title="RT-DETR vs YOLO Vergleich") as demo:
+    gr.Markdown("# 🔍 Objekterkennung: RT-DETR (80 COCO) vs YOLO11n (365 Objects365)")
+    gr.Markdown("Vergleiche RT-DETR mit Standard COCO gegen YOLO mit erweitertem Objects365 Datensatz")
     
     with gr.Row():
         with gr.Column():
@@ -208,8 +245,8 @@ with gr.Blocks(title="RT-DETR Modellvergleich") as demo:
             detect_btn = gr.Button("🚀 Modelle vergleichen", variant="primary")
         
     with gr.Row():
-        r50_output = gr.Image(label="RT-DETR R50 (schnell)")
-        r101_output = gr.Image(label="RT-DETR R101 (genau)")
+        rtdetr_output = gr.Image(label="RT-DETR R50 (80 COCO Klassen)")
+        yolo_output = gr.Image(label="YOLO11n (365 Objects365 Klassen)")
     
     analysis_output = gr.Markdown(label="Vergleichsanalyse")
     
@@ -217,7 +254,7 @@ with gr.Blocks(title="RT-DETR Modellvergleich") as demo:
     detect_btn.click(
         fn=compare_models,
         inputs=[input_image, confidence_slider],
-        outputs=[r50_output, r101_output, analysis_output]
+        outputs=[rtdetr_output, yolo_output, analysis_output]
     )
     
     # Beispiele
@@ -227,7 +264,7 @@ with gr.Blocks(title="RT-DETR Modellvergleich") as demo:
             ["example2.jpg", 0.3],
         ],
         inputs=[input_image, confidence_slider],
-        outputs=[r50_output, r101_output, analysis_output],
+        outputs=[rtdetr_output, yolo_output, analysis_output],
         fn=compare_models,
         cache_examples=False
     )