23.data_loader_visualizer.py

import cv2
from data_loader import DataLoader
import os
import numpy as np
from typing import List, Dict, Tuple, Optional
import json

class RecordingVisualizer:
    def __init__(self, video_path: str, jsonl_path: str):
        self.cap = cv2.VideoCapture(video_path)
        self.frame_events = self.load_jsonl(jsonl_path)
        self.total_frames = len(self.frame_events)
        
        if self.total_frames == 0:
            raise ValueError("No events found in JSONL file")
            
        # Verify video file is valid
        if not self.cap.isOpened():
            raise ValueError(f"Could not open video file: {video_path}")
            
        # Get frame dimensions
        self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    @staticmethod
    def load_jsonl(jsonl_path: str) -> List[Dict]:
        """Load JSONL file into list of events"""
        events = []
        with open(jsonl_path, 'r') as f:
            for line in f:
                events.append(json.loads(line.strip()))
        return events
    
    def parse_mouse_position(self, mouse_pos: Optional[List[int]]) -> Optional[Tuple[int, int]]:
        """Safely parse mouse position data"""
        if not mouse_pos or not isinstance(mouse_pos, (list, tuple)) or len(mouse_pos) != 2:
            return None
        try:
            x, y = int(mouse_pos[0]), int(mouse_pos[1])
            return (x, y)
        except (ValueError, TypeError):
            return None

    def is_mouse_in_frame(self, mouse_pos: Optional[Tuple[int, int]]) -> bool:
        """Check if mouse coordinates are within frame bounds"""
        if not mouse_pos:
            return False
        x, y = mouse_pos
        return (0 <= x < self.frame_width) and (0 <= y < self.frame_height)

    def get_edge_intersection(self, mouse_pos: Tuple[int, int]) -> Tuple[int, int]:
        """Get the point where a line from the center to mouse_pos intersects frame edge"""
        x, y = mouse_pos
        center_x = self.frame_width // 2
        center_y = self.frame_height // 2
        
        # Vector from center to mouse position
        dx = x - center_x
        dy = y - center_y
        
        # Find intersection with frame edges
        if abs(dx) > abs(dy):
            # Intersects with left or right edge
            if dx > 0:
                edge_x = self.frame_width - 1
            else:
                edge_x = 0
            # Scale y accordingly
            if dx != 0:
                edge_y = int(center_y + dy * (edge_x - center_x) / dx)
                edge_y = min(max(edge_y, 0), self.frame_height - 1)
            else:
                edge_y = center_y
        else:
            # Intersects with top or bottom edge
            if dy > 0:
                edge_y = self.frame_height - 1
            else:
                edge_y = 0
            # Scale x accordingly
            if dy != 0:
                edge_x = int(center_x + dx * (edge_y - center_y) / dy)
                edge_x = min(max(edge_x, 0), self.frame_width - 1)
            else:
                edge_x = center_x
                
        return edge_x, edge_y

    def draw_offscreen_indicator(self, frame: np.ndarray, mouse_pos: Tuple[int, int]) -> np.ndarray:
        """Draw an arrow indicating off-screen mouse position"""
        if not mouse_pos:
            return frame
            
        x, y = mouse_pos
        # Get intersection point with frame edge
        edge_x, edge_y = self.get_edge_intersection(mouse_pos)
        
        # Draw arrow from edge pointing towards off-screen mouse
        arrow_length = 20
        center_x = self.frame_width // 2
        center_y = self.frame_height // 2
        
        # Calculate arrow direction
        dx = x - center_x
        dy = y - center_y
        magnitude = np.sqrt(dx*dx + dy*dy)
        if magnitude > 0:
            dx = dx / magnitude
            dy = dy / magnitude
        
        # Draw arrow
        start_point = (edge_x - int(dx * arrow_length), edge_y - int(dy * arrow_length))
        end_point = (edge_x, edge_y)
        
        cv2.arrowedLine(
            frame,
            start_point,
            end_point,
            (0, 255, 0),
            2,
            tipLength=0.3
        )
        
        # Add distance indicator
        distance = int(np.sqrt((x - edge_x)**2 + (y - edge_y)**2))
        cv2.putText(
            frame,
            f"{distance}px",
            (edge_x + 5, edge_y + 5),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,
            (0, 255, 0),
            1
        )
        
        return frame

    def draw_mouse_status(self, frame: np.ndarray, 
                         raw_mouse_pos: Optional[List[int]], 
                         mouse_buttons: List[str],
                         mouse_scroll: Optional[Tuple[int, int]],
                         y_start: int = 80) -> np.ndarray:
        """Draw mouse events on frame"""
        # Mouse position
        cv2.putText(
            frame,
            "Mouse:",
            (10, y_start),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (255, 255, 255),
            1
        )
        
        # Parse mouse position safely
        mouse_pos = self.parse_mouse_position(raw_mouse_pos)
        
        if mouse_pos:
            x, y = mouse_pos
            pos_text = f"Position: ({x}, {y})"
            
            if self.is_mouse_in_frame(mouse_pos):
                # Draw crosshair for in-frame mouse position
                cv2.drawMarker(
                    frame, 
                    (x, y), 
                    (0, 255, 0), 
                    markerType=cv2.MARKER_CROSS, 
                    markerSize=15, 
                    thickness=1
                )
            else:
                # Draw off-screen indicator
                frame = self.draw_offscreen_indicator(frame, mouse_pos)
                pos_text += " (Off-screen)"
        else:
            pos_text = "Position: None"
            
        cv2.putText(
            frame,
            pos_text,
            (10, y_start + 20),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,
            (255, 255, 255),
            1
        )
        
        # Mouse buttons
        button_text = ", ".join(mouse_buttons) if mouse_buttons else "None"
        cv2.putText(
            frame,
            f"Buttons: {button_text}",
            (10, y_start + 40),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,
            (255, 255, 255),
            1
        )
        
        # Scroll info
        if mouse_scroll:
            scroll_text = f"Scroll: dx={mouse_scroll[0]}, dy={mouse_scroll[1]}"
        else:
            scroll_text = "Scroll: None"
            
        cv2.putText(
            frame,
            scroll_text,
            (10, y_start + 60),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,
            (255, 255, 255),
            1
        )
        
        return frame
    
    def draw_keyboard_status(self, frame: np.ndarray, keys: List[str], y_start: int = 30) -> np.ndarray:
        """Draw keyboard events on frame"""
        cv2.putText(
            frame,
            "Keyboard:",
            (10, y_start),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (255, 255, 255),
            1
        )
        
        # Display active keys
        key_text = ", ".join(keys) if keys else "None"
        cv2.putText(
            frame,
            f"Active Keys: {key_text}",
            (10, y_start + 20),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,
            (255, 255, 255),
            1
        )
        
        return frame
    
    def draw_frame_info(self, frame: np.ndarray, frame_number: int, timestamp: float) -> np.ndarray:
        """Draw frame information"""
        cv2.putText(
            frame,
            f"Frame: {frame_number} | Time: {timestamp:.2f}s",
            (10, 20),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (255, 255, 255),
            1
        )
        return frame
    
    def visualize(self):
        """Display frames with overlaid input information"""
        frame_idx = 0
        
        while frame_idx < self.total_frames:
            ret, frame = self.cap.read()
            if not ret:
                print("Error reading frame")
                break
                
            event = self.frame_events[frame_idx]
            
            # Draw information layers
            frame = self.draw_frame_info(frame, event['frame_number'], event['timestamp'])
            frame = self.draw_keyboard_status(frame, event['keyboard_keys'])
            frame = self.draw_mouse_status(
                frame,
                event['mouse_position'],
                event['mouse_buttons'],
                event['mouse_scroll']
            )
            
            # Display the frame
            cv2.imshow('Recording Playback', frame)
            
            # Handle keyboard input
            key = cv2.waitKey(0)  # Wait for key press
            
            if key == ord('q'):  # Quit
                break
            elif key == ord('a'):  # Previous frame
                frame_idx = max(0, frame_idx - 1)
                self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            elif key == ord('d'):  # Next frame
                frame_idx = min(self.total_frames - 1, frame_idx + 1)
            elif key == ord('s'):  # Play/pause
                while True:
                    ret, frame = self.cap.read()
                    if not ret or frame_idx >= self.total_frames - 1:
                        break
                        
                    frame_idx += 1
                    event = self.frame_events[frame_idx]
                    
                    frame = self.draw_frame_info(frame, event['frame_number'], event['timestamp'])
                    frame = self.draw_keyboard_status(frame, event['keyboard_keys'])
                    frame = self.draw_mouse_status(
                        frame,
                        event['mouse_position'],
                        event['mouse_buttons'],
                        event['mouse_scroll']
                    )
                    
                    cv2.imshow('Recording Playback', frame)
                    if cv2.waitKey(33) == ord('s'):  # ~30 FPS playback
                        break
        
        self.cap.release()
        cv2.destroyAllWindows()
if __name__ == "__main__":
    # Example usage:
    dataset_dir = './recordings/'  # Replace with your dataset directory
    batch_size = 2
    n_workers = 2
    n_epochs = 1
    
    data_loader = DataLoader(dataset_dir, n_workers, batch_size, n_epochs)

    for i, (frames, actions, episode_ids) in enumerate(data_loader):
        print(f"Batch {i+1}:")
        print(f"  - Frames (first frame in batch): shape={frames[0].shape}, type={frames[0].dtype}")
        print(f"  - Actions (first action in batch): {actions[0]}")
        print(f"  - Episode IDs: {episode_ids}")
        
        
        # Display video with annotation
        video_path = os.path.join(dataset_dir, f"recording_20241222_132253_recording.mp4")
        jsonl_path = os.path.join(dataset_dir, f"recording_20241222_132253_events.jsonl")
        
        # load data for display
        recorder_visualizer = RecordingVisualizer(video_path, jsonl_path)
        recorder_visualizer.visualize()
        if i >= 2:
            break