Spaces:
Sleeping
Sleeping
| import os | |
| import glob | |
| import gradio as gr | |
| import shutil | |
| from ultralytics import YOLO | |
| # Initialize YOLO model | |
| model = YOLO('./best.pt') | |
| def check_ffmpeg(): | |
| """Check if ffmpeg is installed.""" | |
| if shutil.which("ffmpeg") is None: | |
| raise EnvironmentError("ffmpeg is not installed or not found in the system PATH. Please install ffmpeg to proceed.") | |
| def process_video(video_path, output_option): | |
| # Check if ffmpeg is installed | |
| check_ffmpeg() | |
| # Run the YOLO model on the video, specifying the tracker configuration | |
| results = model.track(video_path, save=True, tracker="bytetrack.yaml", persist=True) | |
| # Get the video file name without extension and directory | |
| video_name = os.path.splitext(os.path.basename(video_path))[0] | |
| # Find the latest directory created in 'runs/detect/' | |
| output_dir = max(glob.glob('./runs/detect/*'), key=os.path.getmtime) | |
| # Find the saved video file in the latest directory with the specific video name | |
| video_files = glob.glob(os.path.join(output_dir, f'{video_name}*.avi')) | |
| if not video_files: | |
| raise Exception(f"No .avi video files found in directory {output_dir} for {video_name}") | |
| # Convert the video to mp4 | |
| mp4_files = [] | |
| for file in video_files: | |
| mp4_file = f"{file[:-4]}.mp4" | |
| # Quote the file paths to handle spaces and overwrite without prompt | |
| os.system(f"ffmpeg -y -i \"{file}\" -vcodec libx264 \"{mp4_file}\"") | |
| os.remove(file) # Remove the original .avi file after conversion | |
| mp4_files.append(mp4_file) | |
| # Check again for the converted .mp4 video files specifically for the processed video | |
| matched_mp4_files = [file for file in mp4_files if video_name in file] | |
| if not matched_mp4_files: | |
| raise Exception(f"No .mp4 video files found in directory {output_dir} after conversion for {video_name}.") | |
| # Initialize object count | |
| object_count = 0 | |
| # Calculate object count based on results | |
| for result in results: | |
| if hasattr(result, 'boxes'): | |
| object_count += len(result.boxes) | |
| # Set default outputs | |
| video_output = matched_mp4_files[0] if matched_mp4_files else None | |
| count_output = object_count | |
| # Determine the output based on user's choice | |
| if output_option == "Count": | |
| return count_output, None | |
| elif output_option == "Video": | |
| return None, video_output | |
| elif output_option == "Both": | |
| return count_output, video_output | |
| # Define Gradio inputs | |
| video_input = gr.Video() | |
| output_option_input = gr.Radio(choices=["Count", "Video", "Both"], label="Select output type") | |
| # Define examples for the interface | |
| examples = [ | |
| [os.path.abspath("example_video1.mp4"), "Count"], | |
| [os.path.abspath("example_video2.mp4"), "Video"], | |
| [os.path.abspath("example_video3.mp4"), "Both"] | |
| ] | |
| video_interface = gr.Interface( | |
| fn=process_video, | |
| inputs=[video_input, output_option_input], | |
| outputs=[gr.Textbox(label="Object Count"), "video"], # Ensure two outputs are defined | |
| title="YOLO Video Tracking Application", | |
| description="A simple application to track objects in a video using YOLO model. Upload your own video, or click one of the examples to load them.", | |
| article="""<div> | |
| <p style="text-align: center">Upload a video file and select the type of output you want: object count, processed video, or both. Then, hit submit to process the video.</p> | |
| </div>""", | |
| examples=examples | |
| ) | |
| # Deploy the interface with share enabled | |
| gr.TabbedInterface([video_interface], ["Track Video"]).launch(share=True) | |