Update app.py

app.py

@@ -5,62 +5,90 @@ from huggingface_hub import snapshot_download
 import os
 import tempfile
 import cv2
+import zipfile
+import shutil
 
 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_image(pilimg, conf_threshold, iou_threshold):
-    result = detection_model.predict(pilimg, conf=conf_threshold, iou=iou_threshold)
-    img_bgr = result[0].plot()
-    out_img = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB for PIL
+def predict(files, conf_threshold, iou_threshold):
+    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")
+            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
+        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))
 
-    # Save to temp file and return path
-    tmp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-    out_img.save(tmp_img.name)
-    return tmp_img.name
+            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)
 
-def predict_video(video_path, conf_threshold, iou_threshold):
-    cap = cv2.VideoCapture(video_path)
-    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 out_path
+        else:
+            return "Unsupported file type."
+    else:
+        output_dir = tempfile.mkdtemp()
+        annotated_images = []
 
-    cap.release()
-    out_writer.release()
-    return out_path
+        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}")
 
-def predict(file, conf_threshold, iou_threshold):
-    ext = os.path.splitext(file.name)[1].lower()
-    if ext in ['.jpg', '.jpeg', '.png']:
-        img = Image.open(file).convert("RGB")
-        return predict_image(img, conf_threshold, iou_threshold)
-    elif ext in ['.mp4', '.mov', '.avi']:
-        return predict_video(file.name, conf_threshold, iou_threshold)
-    else:
-        return "Unsupported file type. Please upload an image or video."
+        zip_path = shutil.make_archive(output_dir, 'zip', output_dir)
+        return annotated_images, zip_path
 
 REPO_ID = "sensura/belisha-beacon-zebra-crossing-yoloV8"
 detection_model = load_model(REPO_ID)
 
-gr.Interface(
+def dynamic_output(file_list):
+    if len(file_list) == 1:
+        ext = os.path.splitext(file_list[0].name)[1].lower()
+        if ext in ['.jpg', '.jpeg', '.png']:
+            return gr.Image(type="filepath", label="Annotated Image")
+        elif ext in ['.mp4', '.mov', '.avi']:
+            return gr.Video(label="Annotated Video")
+    else:
+        return [
+            gr.Gallery(label="Annotated Images").style(grid=3, height="auto"),
+            gr.File(label="Download All Annotated Images (ZIP)")
+        ]
+
+interface = gr.Interface(
     fn=predict,
     inputs=[
-        gr.File(label="Upload Image or Video"),
+        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")
     ],
-    outputs=gr.File(label="Processed Output")
-).launch(share=True)
+    outputs=dynamic_output,
+    live=False
+)
+
+interface.launch(share=True)