app update

app.py

@@ -1,10 +1,8 @@
+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 tempfile
 
 #public model path location
 #MODEL_REPO_ID = "mintheinwin/3907578Y"
@@ -12,93 +10,34 @@ import tempfile
 #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)
-    path = os.path.join(download_dir, "best_int8_openvino_model")
-    detection_model = YOLO(path, task="detect")
+    print(download_dir)
+    path  = os.path.join(download_dir, "best_int8_openvino_model")
+    print(path)
+    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"
 
-#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
+#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
     return out_pilimg
-
-#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("### Detection Results:")
-            output_display = gr.Output(label="Output")
-
-    submit_btn = gr.Button("Detect Objects")
-
-    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."
-
-    submit_btn.click(fn=process_file, inputs=input_file, outputs=output_display)
-
-# Launch app
-interface.launch(share=True)
+    
+#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)