Update app.py

app.py

@@ -1,68 +1,54 @@
+import os
 import gradio as gr
 import cv2
-import requests
-import os
 from ultralytics import YOLO
 
-# Define file URLs for images and videos
-file_urls = [
-    'https://drive.google.com/file/d/1rvuphnn3BV4NdILrQE72jU7fxA79SiYn/view',  # Image
-    'https://drive.google.com/file/d/16gu9cLamGFrM5DRd1WJyk_6Xt9v0S7go/view',  # Image
-    'https://drive.google.com/file/d/1UgZi54js65f5qGhNF3nGLZwIN5nrUek6/view',  # Video
-]
+# Define the folder containing the images and video
+folder_path = "info"  # Replace with your folder name or path
+
+# Get list of files from the folder
+image_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
+video_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.lower().endswith(('.mp4', '.avi', '.mov'))]
 
-# 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)
+# Ensure the folder contains the expected number of files
+if len(image_files) < 2 or len(video_files) < 1:
+    raise ValueError("Folder must contain at least 2 images and 1 video.")
 
-# 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")
+# Select the first two images and the first video
+image_examples = [[image_files[0]], [image_files[1]]]
+video_examples = [[video_files[0]]]
 
 # 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
+    annotated_image = results[0].plot()
     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
-
+    out_frames = []
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
     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))
-
+    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
     for frame in out_frames:
         writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-
     writer.release()
     return output_path
 
@@ -77,7 +63,7 @@ interface_image = gr.Interface(
     examples=image_examples,
 )
 
-inputs_video = gr.Video(label="Input Video")  # Removed type argument
+inputs_video = gr.Video(label="Input Video")
 outputs_video = gr.Video(label="Annotated Output")
 interface_video = gr.Interface(
     fn=show_preds_video,
@@ -91,4 +77,4 @@ interface_video = gr.Interface(
 gr.TabbedInterface(
     [interface_image, interface_video],
     tab_names=['Image Inference', 'Video Inference']
-).launch(share=True)
+).launch(share=True)