Khalida1w commited on
Commit
1d9c414
1 Parent(s): 47e9e2b

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +63 -0
  2. requirements.txt +2 -0
app.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import glob
3
+ import gradio as gr
4
+ import shutil
5
+ from ultralytics import YOLO
6
+
7
+ # Initialize YOLO model
8
+ model = YOLO('./best.pt')
9
+
10
+ def check_ffmpeg():
11
+ """Check if ffmpeg is installed."""
12
+ if shutil.which("ffmpeg") is None:
13
+ raise EnvironmentError("ffmpeg is not installed or not found in the system PATH. Please install ffmpeg to proceed.")
14
+
15
+ def process_video(video_path):
16
+ # Check if ffmpeg is installed
17
+ check_ffmpeg()
18
+
19
+ # Run the YOLO model on the video, specifying the tracker configuration
20
+ results = model.track(video_path, save=True, tracker="bytetrack.yaml", persist=True)
21
+
22
+ # Get the video file name without extension and directory
23
+ video_name = os.path.splitext(os.path.basename(video_path))[0]
24
+
25
+ # Find the latest directory created in 'runs/detect/'
26
+ output_dir = max(glob.glob('./runs/detect/*'), key=os.path.getmtime)
27
+
28
+ # Find the saved video file in the latest directory with the specific video name
29
+ video_files = glob.glob(os.path.join(output_dir, f'{video_name}*.avi'))
30
+
31
+ if not video_files:
32
+ raise Exception(f"No .avi video files found in directory {output_dir} for {video_name}")
33
+
34
+ # Convert the video to mp4
35
+ mp4_files = []
36
+ for file in video_files:
37
+ mp4_file = f"{file[:-4]}.mp4"
38
+ # Quote the file paths to handle spaces and overwrite without prompt
39
+ os.system(f"ffmpeg -y -i \"{file}\" -vcodec libx264 \"{mp4_file}\"")
40
+ os.remove(file) # Remove the original .avi file after conversion
41
+ mp4_files.append(mp4_file)
42
+
43
+ # Check again for the converted .mp4 video files specifically for the processed video
44
+ matched_mp4_files = [file for file in mp4_files if video_name in file]
45
+
46
+ if not matched_mp4_files:
47
+ raise Exception(f"No .mp4 video files found in directory {output_dir} after conversion for {video_name}.")
48
+
49
+ # Return the first video file found (should correspond to the latest processed video)
50
+ return matched_mp4_files[0]
51
+
52
+ # Define Gradio inputs
53
+ video_input = gr.Video()
54
+
55
+ video_interface = gr.Interface(
56
+ fn=process_video,
57
+ inputs=video_input,
58
+ outputs="video", # Display the video output
59
+ description="YOLO Video Tracking"
60
+ )
61
+
62
+ # Deploy the interface with share enabled
63
+ gr.TabbedInterface([video_interface], ["Track Video"]).launch(share=True)
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ ultralytics==8.2.82
2
+ lapx==0.5.10.post1