Spaces:
Build error
Build error
| import os | |
| import csv | |
| import shutil # Added to copy files | |
| from dataclasses import dataclass | |
| import model.video_classifier as classifier | |
| device = 'cpu' | |
| model_path = './model/best.pt' | |
| video_folder = './videos' | |
| output_folder = './results' | |
| # Ensure the directories exist | |
| os.makedirs(video_folder, exist_ok=True) | |
| os.makedirs(output_folder, exist_ok=True) | |
| model = classifier.get_model(model_path) | |
| class VideoResult: | |
| video_path: str | |
| label_counts: dict | |
| crash_events: list | |
| video_name: str | |
| is_simulation: int | |
| def result_csv(student, simulation_results, real_results): | |
| results = simulation_results + real_results | |
| file_name = 'crash_results.csv' | |
| csv_file = os.path.join(output_folder, file_name) | |
| csv_columns = [ | |
| 'Student', 'VideoFileName', 'IsSimulation', 'UniqueCrashes', 'CrashFrames', 'FlightFrames', | |
| 'NoDroneFrames', 'NoSignalFrames', 'NoStartedFrames', 'StartedFrames', 'UnstableFrames', | |
| 'LandingFrames', 'CrashTimes', 'CrashDurations' | |
| ] | |
| with open(csv_file, 'w', newline='') as csvfile: | |
| writer = csv.DictWriter(csvfile, fieldnames=csv_columns) | |
| writer.writeheader() | |
| for result in results: | |
| unique_crashes = len(result.crash_events) | |
| crash_times = '' | |
| crash_durations = '' | |
| if unique_crashes > 0: | |
| for j, (start, end) in enumerate(result.crash_events, start=1): | |
| start_formatted = classifier.format_time(start) | |
| end_formatted = classifier.format_time(end) | |
| duration = end - start | |
| crash_times += f"[{start_formatted},{end_formatted}]" | |
| crash_durations += f"{duration:.2f}" | |
| if unique_crashes - j != 0: | |
| crash_times += "," | |
| crash_durations += "," | |
| else: | |
| crash_times = None | |
| crash_durations = None | |
| writer.writerow({ | |
| 'Student': student, | |
| 'VideoFileName': result.video_name, | |
| 'IsSimulation': result.is_simulation, | |
| 'UniqueCrashes': unique_crashes, | |
| 'CrashFrames': result.label_counts.get("Crash", 0), | |
| 'FlightFrames': result.label_counts.get("Flight", 0), | |
| 'NoDroneFrames': result.label_counts.get("No drone", 0), | |
| 'NoSignalFrames': result.label_counts.get("No signal", 0), | |
| 'NoStartedFrames': result.label_counts.get("No started", 0), | |
| 'StartedFrames': result.label_counts.get("Started", 0), | |
| 'UnstableFrames': result.label_counts.get("Unstable", 0), | |
| 'LandingFrames': result.label_counts.get("Landing", 0), | |
| 'CrashTimes': crash_times, | |
| 'CrashDurations': crash_durations | |
| }) | |
| return csv_file | |
| def remove_videos_and_csv(folder_path): | |
| video_extensions = {'.mp4'} | |
| csv_extension = '.csv' | |
| for filename in os.listdir(folder_path): | |
| file_path = os.path.join(folder_path, filename) | |
| if os.path.isfile(file_path): | |
| _, extension = os.path.splitext(filename) | |
| if extension.lower() in video_extensions or extension.lower() == csv_extension: | |
| os.remove(file_path) | |
| def delete_uploaded_videos(simulation_paths, real_paths): | |
| all_video_paths = simulation_paths + real_paths | |
| for video_file in all_video_paths: | |
| if os.path.exists(video_file): | |
| os.remove(video_file) | |
| def backend(student_name_tb, simulation_files, real_files, frame_output, text_output, file_output, video_output): | |
| yield frame_output, "Starting process...", file_output, video_output | |
| remove_videos_and_csv(output_folder) | |
| yield frame_output, "Processing Input...", file_output, video_output | |
| name = student_name_tb.strip() | |
| simulation_paths = [] | |
| real_paths = [] | |
| # Ensure the video_folder exists | |
| os.makedirs(video_folder, exist_ok=True) | |
| # Store original filenames | |
| simulation_filenames = [os.path.basename(file_path) for file_path in simulation_files] | |
| real_filenames = [os.path.basename(file_path) for file_path in real_files] | |
| # Save the uploaded simulation videos to the 'video_folder' | |
| for i, file_path in enumerate(simulation_files, start=1): | |
| ext = os.path.splitext(file_path)[1] | |
| sim_video_path = os.path.join(video_folder, f'sim_video{i}{ext}') | |
| shutil.copy(file_path, sim_video_path) | |
| simulation_paths.append(sim_video_path) | |
| # Save the uploaded real videos to the 'video_folder' | |
| for i, file_path in enumerate(real_files, start=1): | |
| ext = os.path.splitext(file_path)[1] | |
| real_video_path = os.path.join(video_folder, f'real_video{i}{ext}') | |
| shutil.copy(file_path, real_video_path) | |
| real_paths.append(real_video_path) | |
| simulation_results = [] | |
| real_results = [] | |
| output_video_paths = [] | |
| # Process simulation videos | |
| for i, path in enumerate(simulation_paths, start=1): | |
| # Setup output paths | |
| crash_video_filename = f'Simulation_Video_{i}_Crashes.mp4' | |
| crash_video_path = os.path.join(output_folder, crash_video_filename) | |
| labeled_video_filename = f'Simulation_Video_{i}_Labeled.mp4' | |
| labeled_video_path = os.path.join(output_folder, labeled_video_filename) | |
| output_video_paths.extend([crash_video_path, labeled_video_path]) | |
| yield frame_output, f"Processing Simulation Video #{i}...", file_output, video_output | |
| video_gen = classifier.video_classification( | |
| path, labeled_video_path, crash_video_path, model, min_crash_duration=2.0) | |
| for item in video_gen: | |
| if item['type'] == 'frame': | |
| frame_out = item['frame'] | |
| progress_text = item['progress_text'] | |
| yield frame_out, progress_text, file_output, video_output | |
| elif item['type'] == 'results': | |
| label_counts = item['label_counts'] | |
| crash_events = item['crash_events'] | |
| simulation_results.append(VideoResult( | |
| video_path=path, | |
| label_counts=label_counts, | |
| crash_events=crash_events, | |
| video_name=simulation_filenames[i - 1], | |
| is_simulation=1)) | |
| break | |
| # Process real videos | |
| for i, path in enumerate(real_paths, start=1): | |
| crash_video_filename = f'Real_Video_{i}_Crashes.mp4' | |
| crash_video_path = os.path.join(output_folder, crash_video_filename) | |
| labeled_video_filename = f'Real_Video_{i}_Labeled.mp4' | |
| labeled_video_path = os.path.join(output_folder, labeled_video_filename) | |
| output_video_paths.extend([crash_video_path, labeled_video_path]) | |
| yield frame_output, f"Processing Real Video #{i}...", file_output, video_output | |
| video_gen = classifier.video_classification( | |
| path, labeled_video_path, crash_video_path, model, min_crash_duration=2.0) | |
| for item in video_gen: | |
| if item['type'] == 'frame': | |
| frame_out = item['frame'] | |
| progress_text = item['progress_text'] | |
| yield frame_out, progress_text, file_output, video_output | |
| elif item['type'] == 'results': | |
| label_counts = item['label_counts'] | |
| crash_events = item['crash_events'] | |
| real_results.append(VideoResult( | |
| video_path=path, | |
| label_counts=label_counts, | |
| crash_events=crash_events, | |
| video_name=real_filenames[i - 1], | |
| is_simulation=0)) | |
| break | |
| yield frame_output, "Outputting Results to CSV...", file_output, video_output | |
| csv_file = result_csv(name, simulation_results, real_results) | |
| yield frame_output, "Finishing process...", file_output, video_output | |
| # Delete the uploaded video files from the server | |
| delete_uploaded_videos(simulation_paths, real_paths) | |
| yield frame_output, "Download Ready.", csv_file, output_video_paths | |