Update app.py

app.py

@@ -20,10 +20,6 @@
 
 # if __name__ == "__main__":
 #     demo.launch()
-
-import gradio as gr
-import torch
-import cv2
 # from IPython.display import clear_output
 # import os, urllib.request
 # import subprocess
@@ -126,33 +122,37 @@ import cv2
 # print('Installation finished.')
 
 # Define the face detector function
+
+
+import gradio as gr
+import torch
+import cv2
+import os
+
+# Load the model
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
+
 def detect_image(image):
     results = model(image)
-    
     return results.render()[0]
-    
-def detect_video(video_path):
 
+def detect_video(video_path):
     video = cv2.VideoCapture(video_path)
-    
     frame_count = 0
     while True:
-    
         success, frame = video.read()
         if not success:
             break
-        frame = model.predict();
+        frame = model(frame)[0].numpy().astype('uint8')
         frame_output_path = f'frame_{frame_count}.jpg'  # Replace with your desired output path
         cv2.imwrite(frame_output_path, frame)
-        
         frame_count += 1
-    
     video.release()
     cv2.destroyAllWindows()
 
-    frame_rate = 30  # Adjust as needed
+    frame_rate = video.get(cv2.CAP_PROP_FPS)
     
-    image_dir = 'path_to_image_directory'  # Replace with the directory containing the image files
+    image_dir = '.'  # Replace with the directory containing the image files (frames)
     image_files = sorted(os.listdir(image_dir))
     
     image_path = os.path.join(image_dir, image_files[0])
@@ -185,7 +185,6 @@ vid_interface = gr.Interface(
     outputs="video",
     title="Video"
 )
-
 # Create a list of interfaces
 interfaces = [img_interface, vid_interface]