Update multi_video_app.py

multi_video_app.py

@@ -3,37 +3,30 @@ warnings.filterwarnings("ignore")
 import gradio as gr
 import re
 from typing import Dict, List
-import csv
 import os
 import torch
+import pandas as pd
 from src.video_model import describe_video 
 from src.utils import parse_string, parse_annotations
 
-
-
-# Function to save data to a CSV file
+# Function to save data to a CSV file using pandas
 def save_to_csv(observations: List[Dict], output_dir: str = "outputs") -> str:
     if not os.path.exists(output_dir):
         os.makedirs(output_dir)
     
+    # Convert the list of dictionaries to a pandas DataFrame
+    df = pd.DataFrame(observations)
+    
+    # Specify the CSV file path
     csv_file = os.path.join(output_dir, "video_observations.csv")
     
-    with open(csv_file, mode='w', newline='') as file:
-        writer = csv.writer(file)
-        writer.writerow(["video_name", "standing", "hands.free", "indoors", "screen.interaction_yes"])
-        for observation in observations:
-            writer.writerow([
-                observation['video_name'], 
-                observation['standing'], 
-                observation['hands.free'], 
-                observation['indoors'], 
-                observation['screen.interaction_yes']
-            ])
+    # Save the DataFrame to a CSV file
+    df.to_csv(csv_file, index=False)
     
     return csv_file
 
 # Function to process a single video and return the observation data
-def process_single_video(video_path: str, standing, hands, location, screen) -> Dict:
+def process_single_video(video_path, standing, hands, location, screen) -> Dict:
     video_name = os.path.basename(video_path)  # Extract video name from the path
     query = "Describe this video in detail and answer the questions"
     additional_info = []
@@ -57,8 +50,8 @@ def process_single_video(video_path: str, standing, hands, location, screen) ->
     final_prompt = final_query + " " + end_query
     
     # Assuming your describe_video function handles the video processing
-    final_response = describe_video(video_path, final_prompt)
-
+    response = describe_video(video_path, final_prompt)
+    final_response = f"<video_name>{video_name}</video_name>" + " \n" + response
 
     conditions = {
         'standing': (standing, 'standing: 1', 'standing: None'),
@@ -71,20 +64,24 @@ def process_single_video(video_path: str, standing, hands, location, screen) ->
         if not condition:
             final_response = final_response.replace(to_replace, replacement)
 
-    return final_response 
-    
+    # Parse the response to extract video name and annotations
+    parsed_content = parse_string(final_response, ["video_name", "annotation"])
+    video_name = parsed_content['video_name'][0] if parsed_content['video_name'] else None
+    annotations_dict = parse_annotations(parsed_content['annotation']) if parsed_content['annotation'] else {}
 
+    # Return the observation as a dictionary
+    return {'video_name': video_name, **annotations_dict}
 
 # Function to process all videos in a folder
-def process_multiple_videos(video_files: List[str], sitting, hands, location, screen):
+def process_multiple_videos(video_files: List[str], standing, hands, location, screen):
     all_observations = []
 
     for video_path in video_files:
-        observation = process_single_video(video_path, sitting, hands, location, screen)
-        if "error" not in observation:
+        observation = process_single_video(video_path, standing, hands, location, screen)
+        if observation['video_name']:  # Only add valid observations
             all_observations.append(observation)
         else:
-            print(observation["error"])  # Log any errors
+            print("Error processing video:", video_path)  # Log any errors
 
         # Clear GPU cache
         torch.cuda.empty_cache()
@@ -94,8 +91,9 @@ def process_multiple_videos(video_files: List[str], sitting, hands, location, sc
     return "Processing completed. Download the CSV file.", csv_file
 
 # Gradio interface
-def gradio_interface(video_files, sitting, hands, location, screen):
-    return process_multiple_videos(video_files, sitting, hands, location, screen)
+def gradio_interface(video_files, standing, hands, location, screen):
+    video_file_paths = [video.name for video in video_files]  # Extract file paths from uploaded files
+    return process_multiple_videos(video_file_paths, standing, hands, location, screen)
 
 # Inputs
 video_files = gr.File(file_count="multiple", file_types=["video"], label="Upload multiple videos")