import os import cv2 import numpy as np import pandas as pd from math import atan, degrees, pi ########## helper functions########## # reading the csv files def read_csv(csv_path): return pd.read_csv(csv_path) # extracting the video frames for which a fixation is detected def extract_frames(video_path, start_frame, end_frame): cap = cv2.VideoCapture(video_path) frames = [] cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) for i in range(start_frame, end_frame + 1): ret, frame = cap.read() if not ret: break frames.append(frame) cap.release() return frames # image stitching algorithm def stitch_images(frames): sift = cv2.SIFT_create() bf = cv2.BFMatcher() img1 = cv2.cvtColor(frames[0], cv2.COLOR_BGR2GRAY) stitched_img = frames[0] for i in range(1, len(frames)): img2 = cv2.cvtColor(frames[i], cv2.COLOR_BGR2GRAY) kp1, des1 = sift.detectAndCompute(img1, None) kp2, des2 = sift.detectAndCompute(img2, None) matches = bf.knnMatch(des1, des2, k=2) good = [] for m in matches: if m[0].distance < 0.5 * m[1].distance: good.append(m[0]) if len(good) >= 4: src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2) dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2) H, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0) height, width, channels = stitched_img.shape stitched_img = cv2.warpPerspective(stitched_img, H, (width, height)) stitched_img[0:frames[i].shape[0], 0:frames[i].shape[1]] = frames[i] img1 = img2 else: raise AssertionError("Can't find enough keypoints.") return stitched_img, H # applying the homography matrix to the fixation points to tranform them to the reference coordinate system of the stitched frame def transform_coordinates(x_mean, y_mean, H): coords = np.array([[x_mean, y_mean]], dtype='float32').reshape(-1, 1, 2) transformed_coords = cv2.perspectiveTransform(coords, H) return transformed_coords[0][0] # walking through the subfolders and preparing the input and output def process_folder(folder_path): csv_files = [f for f in os.listdir(folder_path) if f.startswith('fixation_') and f.endswith('.csv')] video_files = [f for f in os.listdir(folder_path) if f == 'world.mp4'] if not csv_files or not video_files: print(f"No valid CSV or video file found in {folder_path}") return csv_file_path = os.path.join(folder_path, csv_files[0]) video_file_path = os.path.join(folder_path, video_files[0]) output_csv_path = os.path.join(folder_path, 'corrected_' + csv_files[0]) new_output_csv_path = os.path.join(folder_path, 'correct_saccade_' + csv_files[0]) print(f"Processing CSV file: {csv_file_path}") print(f"Processing Video file: {video_file_path}") print(f"Output will be saved to: {output_csv_path}") main(video_file_path, csv_file_path, output_csv_path, output_csv_path, new_output_csv_path) # compute corrected saccade length and azimuth def compute_correct_saccade(corrected_csv_path, output_csv_path): data = pd.read_csv(corrected_csv_path) if len(data) < 2: raise ValueError("The corrected CSV file should have at least two rows.") output_data = [] # Iterate through all consecutive rows for i in range(len(data) - 1): first_row = data.iloc[i] second_row = data.iloc[i + 1] # Calculate length between the transformed coordinates delta_x = second_row['transformed_x'] - first_row['transformed_x'] delta_y = second_row['transformed_y'] - first_row['transformed_y'] length = np.sqrt(delta_x**2 + delta_y**2) # Compute azimuth azimuth = abs(atan(delta_y / delta_x)) * (180 / pi) if delta_y > 0 and delta_x < 0: azimuth = 180 - azimuth elif delta_x > 0 and delta_y < 0: azimuth = 360 - azimuth elif delta_y < 0 and delta_x < 0: azimuth = 180 + azimuth output_data.append({ 'id': first_row['id'], 'first_gaze_timestamp': first_row['time'], 'first_world_index': first_row['world_index'], 'last_gaze_timestamp': second_row['time'], 'last_world_index': second_row['world_index'], 'length': length, 'azimuth': azimuth }) output_df = pd.DataFrame(output_data) output_df.to_csv(output_csv_path, index=False) print(f"Output saved to {output_csv_path}") ########## main function ########## # main function that uses all of the helper functions def main(video_path, csv_path, output_csv_path, corrected_csv_path, saccade_output_csv_path): data = read_csv(csv_path) output_data = [] for index, row in data.iterrows(): start_frame = int(row['start_frame']) end_frame = int(row['end_frame']) frames = extract_frames(video_path, start_frame, end_frame) print(len(frames)) if len(frames) == 1: transformed_coords = (row['x_mean'], row['y_mean']) else: stitched_img, H = stitch_images(frames) print(H) transformed_coords = transform_coordinates(row['x_mean'], row['y_mean'], H) output_data.append({ 'id': row['id'], 'time': row['time'], 'world_index': row['world_index'], 'x_mean': row['x_mean'], 'y_mean': row['y_mean'], 'start_frame': row['start_frame'], 'end_frame': row['end_frame'], 'dispersion': row['dispersion'], 'transformed_x': transformed_coords[0], 'transformed_y': transformed_coords[1] }) # Write to output CSV file in batches of 100 if (index + 1) % 100 == 0: output_df = pd.DataFrame(output_data) if index + 1 == 100: # If it's the first batch output_df.to_csv(output_csv_path, index=False, mode='w') else: output_df.to_csv(output_csv_path, index=False, mode='a', header=False) output_data = [] # Reset the batch # Write any remaining data if output_data: output_df = pd.DataFrame(output_data) if len(data) <= 100: # If the total data is less than or equal to 100 output_df.to_csv(output_csv_path, index=False, mode='w') else: output_df.to_csv(output_csv_path, index=False, mode='a', header=False) print(f"Output saved to {output_csv_path}") compute_correct_saccade(corrected_csv_path, saccade_output_csv_path) # run! if __name__ == "__main__": root_folder = './data' for subdir, dirs, files in os.walk(root_folder): process_folder(subdir)