Update app.py

app.py

@@ -3,15 +3,14 @@ from ultralytics import YOLOWorld
 from PIL import Image
 import cv2
 import pandas as pd
-import numpy as np
-import io
+import tempfile, os
 
 # --- Load model ---
 def load_model():
     last_err = None
     for name in ["yolov8s-worldv2.pt", "yolov8s-world.pt"]:
         try:
-            m = YOLOWorld(name)   # downloads on first run
+            m = YOLOWorld(name)
             return m
         except Exception as e:
             last_err = e
@@ -20,7 +19,7 @@ def load_model():
 
 model = load_model()
 
-# --- Preprocess image ---
+# --- Preprocess ---
 def preprocess(img: Image.Image, max_side=960, do_resize=True):
     img = img.convert("RGB")
     if do_resize:
@@ -29,7 +28,7 @@ def preprocess(img: Image.Image, max_side=960, do_resize=True):
             img.thumbnail((max_side, max_side), Image.LANCZOS)
     return img
 
-# --- Inference function ---
+# --- Inference ---
 def infer(image, prompts, conf, iou, resize):
     image = preprocess(image, do_resize=resize)
     classes = [c.strip() for c in prompts.split(",") if c.strip()]
@@ -40,33 +39,25 @@ def infer(image, prompts, conf, iou, resize):
     results = model(image, conf=conf, iou=iou)[0]
 
     # Annotated image
-    plotted = results.plot()            # np.ndarray in BGR
-    plotted = cv2.cvtColor(plotted, cv2.COLOR_BGR2RGB)  # convert to RGB
+    plotted = results.plot()
+    plotted = cv2.cvtColor(plotted, cv2.COLOR_BGR2RGB)
     annotated_img = Image.fromarray(plotted)
 
-    # Results table
+    # Results table (as list of rows)
     data = []
     for box, cls, conf in zip(results.boxes.xyxy, results.boxes.cls, results.boxes.conf):
         x1, y1, x2, y2 = [int(v) for v in box.tolist()]
-        data.append({
-            "Class": model.names[int(cls)],
-            "Confidence": round(float(conf), 3),
-            "BBox": f"({x1}, {y1}, {x2}, {y2})"
-        })
-    df = pd.DataFrame(data)
+        data.append([model.names[int(cls)], round(float(conf), 3), f"({x1}, {y1}, {x2}, {y2})"])
 
-    # Convert annotated image to downloadable file
-    buf = io.BytesIO()
-    annotated_img.save(buf, format="PNG")
-    buf.seek(0)
+    # Save annotated image to temp file for download
+    tmp_path = os.path.join(tempfile.gettempdir(), "annotated.png")
+    annotated_img.save(tmp_path)
 
-    return annotated_img, df, buf
+    return annotated_img, data, tmp_path
 
-# --- Gradio interface ---
+# --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.Markdown("# 🦉 YOLO-World Demo (Open-Vocabulary Detection on CPU)")
-    gr.Markdown("Upload an image, type objects you want to detect, and see results!")
-
+    gr.Markdown("# 🦉 YOLO-World Demo (CPU • Open-Vocab)")
     with gr.Row():
         with gr.Column():
             image_in = gr.Image(type="pil", label="Upload Image")
@@ -78,7 +69,7 @@ with gr.Blocks() as demo:
 
         with gr.Column():
             img_out = gr.Image(label="Detections")
-            table_out = gr.Dataframe(headers=["Class", "Confidence", "BBox"], label="Results Table")
+            table_out = gr.Dataframe(headers=["Class","Confidence","BBox"], datatype=["str","number","str"])
             file_out = gr.File(label="Download Annotated Image")
 
     run_btn.click(