Update app.py

app.py

@@ -8,87 +8,111 @@ import cv2
 import zipfile
 import shutil
 
+# === Load model ===
 def load_model(repo_id):
     download_dir = snapshot_download(repo_id)
     path = os.path.join(download_dir, "best_int8_openvino_model")
     return YOLO(path, task='detect')
 
-def predict(files, conf_threshold, iou_threshold):
-    image_output = None
-    video_output = None
-    gallery_output = None
-    zip_output = None
+REPO_ID = "sensura/belisha-beacon-zebra-crossing-yoloV8"
+detection_model = load_model(REPO_ID)
+
+# === Single file prediction ===
+def predict_single(file, conf_threshold, iou_threshold):
+    if file is None:
+        return None, None
+
+    ext = os.path.splitext(file.name)[1].lower()
+
+    if ext in ['.jpg', '.jpeg', '.png']:
+        img = Image.open(file).convert("RGB")
+        result = detection_model.predict(img, conf=conf_threshold, iou=iou_threshold)
+        img_bgr = result[0].plot()
+        out_img = Image.fromarray(img_bgr[..., ::-1])
+        tmp_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
+        out_img.save(tmp_path)
+        return tmp_path, None
+
+    elif ext in ['.mp4', '.mov', '.avi']:
+        cap = cv2.VideoCapture(file.name)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+        out_writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            result = detection_model.predict(frame, conf=conf_threshold, iou=iou_threshold)
+            annotated = result[0].plot()
+            out_writer.write(annotated)
+
+        cap.release()
+        out_writer.release()
+        return None, out_path
+
+    else:
+        return None, None
+
+# === Multiple images prediction ===
+def predict_multiple(files, conf_threshold, iou_threshold):
+    if not files:
+        return None, None
 
-    if len(files) == 1:
-        ext = os.path.splitext(files[0].name)[1].lower()
-        if ext in ['.jpg', '.jpeg', '.png']:
-            img = Image.open(files[0]).convert("RGB")
+    output_dir = tempfile.mkdtemp()
+    annotated_images = []
+
+    for file in files:
+        try:
+            img = Image.open(file).convert("RGB")
             result = detection_model.predict(img, conf=conf_threshold, iou=iou_threshold)
             img_bgr = result[0].plot()
             out_img = Image.fromarray(img_bgr[..., ::-1])
-            tmp_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
-            out_img.save(tmp_path)
-            image_output = tmp_path
-        elif ext in ['.mp4', '.mov', '.avi']:
-            cap = cv2.VideoCapture(files[0].name)
-            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            fps = cap.get(cv2.CAP_PROP_FPS)
-            out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-            out_writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
-
-            while cap.isOpened():
-                ret, frame = cap.read()
-                if not ret:
-                    break
-                result = detection_model.predict(frame, conf=conf_threshold, iou=iou_threshold)
-                annotated = result[0].plot()
-                out_writer.write(annotated)
-
-            cap.release()
-            out_writer.release()
-            video_output = out_path
-        else:
-            return "Unsupported file type.", None, None, None
-    else:
-        output_dir = tempfile.mkdtemp()
-        annotated_images = []
-
-        for file in files:
-            try:
-                img = Image.open(file).convert("RGB")
-                result = detection_model.predict(img, conf=conf_threshold, iou=iou_threshold)
-                img_bgr = result[0].plot()
-                out_img = Image.fromarray(img_bgr[..., ::-1])
-                out_path = os.path.join(output_dir, os.path.basename(file.name))
-                out_img.save(out_path)
-                annotated_images.append(out_img)
-            except Exception as e:
-                print(f"Failed to process {file.name}: {e}")
-
-        zip_path = shutil.make_archive(output_dir, 'zip', output_dir)
-        gallery_output = annotated_images
-        zip_output = zip_path
-
-    return image_output, video_output, gallery_output, zip_output
+            out_path = os.path.join(output_dir, os.path.basename(file.name))
+            out_img.save(out_path)
+            annotated_images.append(out_img)
+        except Exception as e:
+            print(f"Failed to process {file.name}: {e}")
 
-REPO_ID = "sensura/belisha-beacon-zebra-crossing-yoloV8"
-detection_model = load_model(REPO_ID)
+    zip_path = shutil.make_archive(output_dir, 'zip', output_dir)
+    return annotated_images, zip_path
+
+# === Gradio Interfaces ===
+
+# Tab 1: Single Image or Video
+single_file_tab = gr.Interface(
+    fn=predict_single,
+    inputs=[
+        gr.File(file_types=["image", "video"], label="Upload Image or Video"),
+        gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Confidence Threshold"),
+        gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="IoU Threshold"),
+    ],
+    outputs=[
+        gr.Image(type="filepath", label="Detected Image"),
+        gr.Video(label="Detected Video")
+    ],
+    title="Detect from a Single Image or Video"
+)
 
-interface = gr.Interface(
-    fn=predict,
+# Tab 2: Multiple Images
+multi_image_tab = gr.Interface(
+    fn=predict_multiple,
     inputs=[
-        gr.File(file_types=["image", "video"], file_count="multiple", label="Upload Image(s) or Video"),
-        gr.Slider(0.1, 1.0, 0.5, step=0.05, label="Confidence Threshold"),
-        gr.Slider(0.1, 1.0, 0.6, step=0.05, label="IoU Threshold")
+        gr.File(file_types=["image"], file_count="multiple", label="Upload Multiple Images"),
+        gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Confidence Threshold"),
+        gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="IoU Threshold"),
     ],
     outputs=[
-        gr.Image(type="filepath", label="Single Image Output"),
-        gr.Video(label="Single Video Output"),
-        gr.Gallery(columns=3, height="auto"),
-        gr.File(label="Download ZIP for Multiple Images")
+        gr.Gallery(label="Detected Gallery", columns=3, height="auto"),
+        gr.File(label="Download Annotated ZIP")
     ],
-    live=False
+    title="Batch Detect from Multiple Images"
 )
 
-interface.launch(share=True)
+# === Tabbed UI Launch ===
+gr.TabbedInterface(
+    [single_file_tab, multi_image_tab],
+    tab_names=["Single File", "Multiple Images"]
+).launch(share=True)