| | import os |
| | import cv2 |
| | import torch |
| | import pandas as pd |
| | from feature_extraction.extractor import VideoFeatureExtractor |
| | from utils.csv_utils import _create_interaction_row |
| |
|
| |
|
| | class VideoDataExtractor: |
| | def __init__(self): |
| | self.extractor = VideoFeatureExtractor() |
| |
|
| | def extract_video_data( |
| | self, |
| | video_path, |
| | output_csv_path=None, |
| | output_folder=None, |
| | show_video=False, |
| | save_video=False, |
| | ): |
| | """ |
| | Extract interaction data from a video file and return a DataFrame. |
| | |
| | Args: |
| | video_path: Path to input video |
| | output_csv_path: Optional path to save CSV output |
| | output_folder: Optional folder to save annotated video |
| | show_video: Whether to display video during processing |
| | save_video: Whether to save output video |
| | |
| | Returns: |
| | pandas.DataFrame containing extracted interactions |
| | """ |
| | cap = None |
| | video_writer = None |
| | csv_data = [] |
| | seen_interactions = set() |
| |
|
| | try: |
| | if not os.path.exists(video_path): |
| | raise FileNotFoundError(f"Video file not found: {video_path}") |
| |
|
| | cap = cv2.VideoCapture(video_path) |
| | if not cap.isOpened(): |
| | raise ValueError("Error: Could not open video file") |
| |
|
| | fps = cap.get(cv2.CAP_PROP_FPS) |
| | frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| | frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| | total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| |
|
| | video_name = os.path.splitext(os.path.basename(video_path))[0] |
| |
|
| | |
| | batch_size, frame_skip = self.extractor.preprocessor.set_resolution_config( |
| | frame_width, frame_height |
| | ) |
| | self.extractor.preprocessor.frame_skip = frame_skip |
| |
|
| | |
| | if output_folder and save_video: |
| | os.makedirs(output_folder, exist_ok=True) |
| | output_video_path = os.path.join( |
| | output_folder, f"{video_name}_detections.mp4" |
| | ) |
| | video_writer = cv2.VideoWriter( |
| | output_video_path, |
| | cv2.VideoWriter_fourcc(*"mp4v"), |
| | fps / frame_skip, |
| | (frame_width, frame_height), |
| | ) |
| |
|
| | |
| | self.extractor.reset() |
| |
|
| | |
| | for frame_idx in range(0, total_frames, frame_skip): |
| | cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) |
| | ret, frame = cap.read() |
| | if not ret: |
| | break |
| |
|
| | frame_data, annotated_frame = self.extractor.extract_features( |
| | frame, frame_idx |
| | ) |
| |
|
| | if frame_data is not None: |
| | for interaction in frame_data["interactions"]: |
| | interaction_id = ( |
| | interaction["person1_id"], |
| | interaction["person2_id"], |
| | frame_idx, |
| | ) |
| |
|
| | if interaction_id not in seen_interactions: |
| | seen_interactions.add(interaction_id) |
| | row = _create_interaction_row( |
| | video_name, |
| | frame_data, |
| | interaction, |
| | frame_width, |
| | frame_height, |
| | ) |
| | csv_data.append(row) |
| |
|
| | if video_writer is not None and annotated_frame is not None: |
| | video_writer.write(annotated_frame) |
| |
|
| | if show_video and annotated_frame is not None: |
| | cv2.imshow("Video Data Extraction", annotated_frame) |
| | if cv2.waitKey(1) & 0xFF == ord("q"): |
| | break |
| |
|
| | if frame_idx % 100 == 0: |
| | torch.cuda.empty_cache() |
| |
|
| | |
| | if csv_data: |
| | df = pd.DataFrame(csv_data) |
| | if output_csv_path: |
| | df.to_csv( |
| | output_csv_path, |
| | mode="a" if os.path.exists(output_csv_path) else "w", |
| | header=not os.path.exists(output_csv_path), |
| | index=False, |
| | ) |
| | return df |
| | else: |
| | return pd.DataFrame() |
| |
|
| | finally: |
| | if cap is not None: |
| | cap.release() |
| | if video_writer is not None: |
| | video_writer.release() |
| | cv2.destroyAllWindows() |
| | torch.cuda.empty_cache() |
| |
|
