Update app.py

app.py

@@ -1,7 +1,94 @@
 import gradio as gr
+import cv2
+import requests
+import os
+from ultralytics import YOLO
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Define file URLs for images and videos
+file_urls = [
+    'https://drive.google.com/file/d/1rvuphnn3BV4NdILrQE72jU7fxA79SiYn/view?usp=sharing',  # Image
+    'https://drive.google.com/file/d/16gu9cLamGFrM5DRd1WJyk_6Xt9v0S7go/view?usp=sharing',  # Image
+    'https://drive.google.com/file/d/1UgZi54js65f5qGhNF3nGLZwIN5nrUek6/view?usp=sharing',  # Video
+]
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# Helper function to download files
+def download_file(url, save_name):
+    if not os.path.exists(save_name):
+        file = requests.get(url)
+        open(save_name, 'wb').write(file.content)
+
+# Download example files
+for i, url in enumerate(file_urls):
+    if url.endswith(".mp4"):
+        download_file(url, f"video.mp4")
+    else:
+        download_file(url, f"image_{i}.jpg")
+
+# Load the YOLO model
+model = YOLO('best.pt')
+
+# Define example paths for Gradio
+image_examples = [["image_0.jpg"], ["image_1.jpg"]]
+video_examples = [["video.mp4"]]
+
+# Function for processing images
+def show_preds_image(image_path):
+    image = cv2.imread(image_path)
+    results = model.predict(source=image_path)
+    annotated_image = results[0].plot()  # YOLO provides a built-in plot function
+    return cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
+
+# Function for processing videos
+def show_preds_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    out_frames = []  # List to store annotated frames
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        results = model.predict(source=frame)
+        annotated_frame = results[0].plot()
+        out_frames.append(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
+
+    cap.release()
+    
+    # Save the annotated video
+    output_path = "annotated_video.mp4"
+    height, width, _ = out_frames[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    writer = cv2.VideoWriter(output_path, fourcc, 20, (width, height))
+
+    for frame in out_frames:
+        writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+
+    writer.release()
+    return output_path
+
+# Gradio interfaces
+inputs_image = gr.Image(type="filepath", label="Input Image")
+outputs_image = gr.Image(type="numpy", label="Output Image")
+interface_image = gr.Interface(
+    fn=show_preds_image,
+    inputs=inputs_image,
+    outputs=outputs_image,
+    title="Tiger & Ibex Detector - Image",
+    examples=image_examples,
+)
+
+inputs_video = gr.Video(label="Input Video")  # Removed type argument
+outputs_video = gr.Video(label="Annotated Output")
+interface_video = gr.Interface(
+    fn=show_preds_video,
+    inputs=inputs_video,
+    outputs=outputs_video,
+    title="Tiger & Ibex Detector - Video",
+    examples=video_examples,
+)
+
+# Combine into a tabbed interface
+gr.TabbedInterface(
+    [interface_image, interface_video],
+    tab_names=['Image Inference', 'Video Inference']
+).launch(share=True)