update app

app.py

@@ -1,8 +1,11 @@
-from ultralytics import YOLO
-from PIL import Image
 import gradio as gr
 from huggingface_hub import snapshot_download
+from ultralytics import YOLO
 import os
+from PIL import Image
+import cv2
+import numpy as np
+import tempfile
 
 #public model path location
 #MODEL_REPO_ID = "mintheinwin/3907578Y"
@@ -10,34 +13,98 @@ import os
 #Organizations model path location
 MODEL_REPO_ID = "ITI107-2024S2/3907578Y"
 
-#load model
+# Load model
 def load_model(repo_id):
     download_dir = snapshot_download(repo_id)
-    print(download_dir)
-    path  = os.path.join(download_dir, "best_int8_openvino_model")
-    print(path)
-    detection_model = YOLO(path, task='detect')
+    path = os.path.join(download_dir, "best_int8_openvino_model")
+    detection_model = YOLO(path, task="detect")
     return detection_model
-    
+
+
 detection_model = load_model(MODEL_REPO_ID)
 
 #Student ID
 student_info = "Student Id: 3907578Y, Name: Min Thein Win"
 
-#prdeict 
-def predict(pilimg):
-    source = pilimg
-    result = detection_model.predict(source, conf=0.5, iou=0.5)
-    img_bgr = result[0].plot()
-    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
+#Prediction for images
+def predict_image(pil_img):
+    result = detection_model.predict(pil_img, conf=0.5, iou=0.5)
+    img_bgr = result[0].plot()  # Annotated image
+    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert to RGB PIL image
     return out_pilimg
-    
-#UI interface
-gr.Markdown("# Wild Animal Detection (Tiger/Lion)")
-gr.Markdown(student_info)
-gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil",label="Input"),
-             outputs=gr.Image(type="pil",label="Output"),
-             title="Wild Animal Detection (Tiger/Lion)",
-             description=student_info,
-             ).launch(share=True)
+
+#Prediction for videos
+def predict_video(video):
+    cap = cv2.VideoCapture(video)
+    frames = []
+    temp_dir = tempfile.mkdtemp()
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Detection
+        result = detection_model.predict(frame, conf=0.5, iou=0.5)
+        annotated_frame = result[0].plot()
+        frames.append(annotated_frame)
+
+    cap.release()
+
+    # Save annotated video
+    height, width, _ = frames[0].shape
+    output_path = os.path.join(temp_dir, "annotated_video.mp4")
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), 20, (width, height))
+
+    for frame in frames:
+        out.write(frame)
+
+    out.release()
+    return output_path
+
+# Unified prediction function
+def unified_predict(file):
+    if isinstance(file, Image.Image):
+        # If the input is a PIL Image, treat it as an image
+        return predict_image(file)
+    elif isinstance(file, str) and file.endswith(('.mp4', '.avi', '.mov')):
+        # If the input is a video file path, treat it as a video
+        return predict_video(file)
+    else:
+        raise ValueError("Unsupported file type. Please upload an image or a video.")
+
+# UI Interface
+with gr.Blocks() as interface:
+    gr.Markdown("# Wild Animal Detection (Tiger/Lion)")
+    gr.Markdown(student_info)
+
+    # Unified Section
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Upload an Image or Video:")
+            input_file = gr.File(label="Input File")
+
+        with gr.Column():
+            gr.Markdown("### Output Results:")
+            output_display = gr.Output(label="Output")
+            
+    clear_btn= gr.Button("CLEAR")
+    submit_btn = gr.Button("SUBMIT")
+
+    def process_file(file):
+        if file.name.endswith((".jpg", ".jpeg", ".png")):
+            pil_image = Image.open(file.name)
+            return predict_image(pil_image)
+        elif file.name.endswith((".mp4", ".avi", ".mov")):
+            return predict_video(file.name)
+        else:
+            return "Unsupported file type. Please upload an image or a video."
+
+    def clear_all():
+        return None, ""
+        
+    submit_btn.click(fn=process_file, inputs=input_file, outputs=output_display)
+    clear_btn.click(fn=clear_all, inputs=None, outputs=[input_file, output_display])
+
+# Launch app
+interface.launch(share=True)