Update app.py

app.py

@@ -1,54 +1,19 @@
 import torch
 from ultralytics import YOLO
 import gradio as gr
-import easyocr
-import numpy as np
-from PIL import Image
-import ultralytics.nn.tasks as tasks
 
-# ✅ Allowlist the YOLOv8 DetectionModel class so PyTorch 2.6 can unpickle it safely
-torch.serialization.add_safe_globals([tasks.DetectionModel])
+# Trusted YOLO weights path
+WEIGHTS_PATH = "/app/yolov8n.pt"
 
-# Load YOLO model (trusted official weights from Ultralytics)
-WEIGHTS_PATH = "yolov8n.pt"  # Should be in your Space folder or auto-downloaded
-model = YOLO(WEIGHTS_PATH, task="detect")
+# Allow loading full model (Option 1 from error)
+torch.serialization.add_safe_globals([YOLO])
 
-# Initialize EasyOCR
-reader = easyocr.Reader(['en'], gpu=False)
+model = YOLO(WEIGHTS_PATH, task="detect")
 
-def detect_and_read(image):
-    # Run YOLO detection
+def detect_objects(image):
     results = model(image)
-    detections = results[0].boxes.xyxy.cpu().numpy()  # x1, y1, x2, y2
-    labels = results[0].names
-
-    output_data = []
-    img_pil = Image.fromarray(image)
-
-    for i, box in enumerate(detections):
-        x1, y1, x2, y2 = [int(v) for v in box[:4]]
-        cropped_img = img_pil.crop((x1, y1, x2, y2))
-        
-        # OCR
-        ocr_result = reader.readtext(np.array(cropped_img))
-        text = " ".join([res[1] for res in ocr_result])
-        
-        output_data.append({
-            "bbox": (x1, y1, x2, y2),
-            "label": labels[int(results[0].boxes.cls[i])],
-            "text": text
-        })
-    
-    return output_data
-
-# Gradio UI
-demo = gr.Interface(
-    fn=detect_and_read,
-    inputs=gr.Image(type="numpy"),
-    outputs="json",
-    title="YOLOv8 + EasyOCR",
-    description="Object detection with YOLOv8 and OCR with EasyOCR"
-)
+    annotated = results[0].plot()
+    return annotated
 
-if __name__ == "__main__":
-    demo.launch()
+demo = gr.Interface(fn=detect_objects, inputs="image", outputs="image")
+demo.launch(server_name="0.0.0.0", server_port=7860)