inital commit

utils/__init__.py

@@ -0,0 +1 @@
+# Utility modules for the YOLO + SAM2 processing pipeline

utils/file_utils.py

@@ -0,0 +1,168 @@
+"""
+File system utilities for the YOLO + SAM2 video processing pipeline.
+"""
+
+import os
+import shutil
+import glob
+from typing import List, Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+def ensure_directory(path: str) -> str:
+    """
+    Ensure directory exists, create if it doesn't.
+    
+    Args:
+        path: Directory path to create
+        
+    Returns:
+        The created directory path
+    """
+    os.makedirs(path, exist_ok=True)
+    logger.debug(f"Ensured directory exists: {path}")
+    return path
+
+def cleanup_directory(path: str, pattern: str = "*") -> int:
+    """
+    Clean up files matching pattern in directory.
+    
+    Args:
+        path: Directory path to clean
+        pattern: File pattern to match (default: all files)
+        
+    Returns:
+        Number of files removed
+    """
+    if not os.path.exists(path):
+        return 0
+    
+    files_to_remove = glob.glob(os.path.join(path, pattern))
+    removed_count = 0
+    
+    for file_path in files_to_remove:
+        try:
+            if os.path.isfile(file_path):
+                os.remove(file_path)
+                removed_count += 1
+            elif os.path.isdir(file_path):
+                shutil.rmtree(file_path)
+                removed_count += 1
+        except OSError as e:
+            logger.warning(f"Failed to remove {file_path}: {e}")
+    
+    if removed_count > 0:
+        logger.info(f"Cleaned up {removed_count} files/directories from {path}")
+    
+    return removed_count
+
+def get_segments_directories(base_dir: str) -> List[str]:
+    """
+    Get list of segment directories sorted by segment number.
+    
+    Args:
+        base_dir: Base directory containing segments
+        
+    Returns:
+        Sorted list of segment directory names
+    """
+    if not os.path.exists(base_dir):
+        return []
+    
+    segments = [d for d in os.listdir(base_dir) 
+                if os.path.isdir(os.path.join(base_dir, d)) and d.startswith("segment_")]
+    
+    # Sort by segment number
+    segments.sort(key=lambda x: int(x.split("_")[1]))
+    
+    logger.debug(f"Found {len(segments)} segment directories in {base_dir}")
+    return segments
+
+def get_video_file_name(segment_index: int) -> str:
+    """
+    Get standardized video filename for a segment.
+    
+    Args:
+        segment_index: Index of the segment
+        
+    Returns:
+        Formatted filename
+    """
+    return f"segment_{str(segment_index).zfill(3)}.mp4"
+
+def file_exists(file_path: str) -> bool:
+    """
+    Check if file exists and is readable.
+    
+    Args:
+        file_path: Path to file
+        
+    Returns:
+        True if file exists and is readable
+    """
+    return os.path.isfile(file_path) and os.access(file_path, os.R_OK)
+
+def create_file_list(segments_dir: str, output_path: str) -> str:
+    """
+    Create ffmpeg-compatible file list for concatenation.
+    
+    Args:
+        segments_dir: Directory containing segment subdirectories
+        output_path: Path to write the file list
+        
+    Returns:
+        Path to the created file list
+    """
+    segments = get_segments_directories(segments_dir)
+    
+    with open(output_path, 'w') as f:
+        for i, segment in enumerate(segments):
+            segment_dir = os.path.join(segments_dir, segment)  
+            output_video = os.path.join(segment_dir, f"output_{i}.mp4")
+            
+            if file_exists(output_video):
+                # Use relative path for ffmpeg
+                relative_path = os.path.relpath(output_video, os.path.dirname(output_path))
+                f.write(f"file '{relative_path}'\\n")
+    
+    logger.info(f"Created file list at {output_path}")
+    return output_path
+
+def safe_remove_file(file_path: str) -> bool:
+    """
+    Safely remove a file with error handling.
+    
+    Args:
+        file_path: Path to file to remove
+        
+    Returns:
+        True if file was removed successfully
+    """
+    try:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+            logger.debug(f"Removed file: {file_path}")
+            return True
+        return False
+    except OSError as e:
+        logger.warning(f"Failed to remove {file_path}: {e}")
+        return False
+
+def get_file_size_mb(file_path: str) -> float:
+    """
+    Get file size in megabytes.
+    
+    Args:
+        file_path: Path to file
+        
+    Returns:
+        File size in MB, or 0 if file doesn't exist
+    """
+    try:
+        if os.path.exists(file_path):
+            size_bytes = os.path.getsize(file_path)
+            return size_bytes / (1024 * 1024)
+        return 0.0
+    except OSError:
+        return 0.0

utils/logging_utils.py

@@ -0,0 +1,52 @@
+"""
+Logging utilities for the YOLO + SAM2 video processing pipeline.
+"""
+
+import logging
+import sys
+from typing import Optional
+
+def setup_logging(level: str = "INFO", log_file: Optional[str] = None):
+    """
+    Setup logging configuration.
+    
+    Args:
+        level: Logging level (DEBUG, INFO, WARNING, ERROR)
+        log_file: Optional log file path
+    """
+    # Convert string level to logging constant
+    numeric_level = getattr(logging, level.upper(), None)
+    if not isinstance(numeric_level, int):
+        raise ValueError(f'Invalid log level: {level}')
+    
+    # Create formatter
+    formatter = logging.Formatter(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    )
+    
+    # Setup console handler
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(numeric_level)
+    console_handler.setFormatter(formatter)
+    
+    # Setup root logger
+    root_logger = logging.getLogger()
+    root_logger.setLevel(numeric_level)
+    root_logger.addHandler(console_handler)
+    
+    # Setup file handler if specified
+    if log_file:
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setLevel(numeric_level)
+        file_handler.setFormatter(formatter)
+        root_logger.addHandler(file_handler)
+    
+    # Reduce noise from some libraries
+    logging.getLogger('ultralytics').setLevel(logging.WARNING)
+    logging.getLogger('matplotlib').setLevel(logging.WARNING)
+    
+    logging.info(f"Logging setup complete - Level: {level}")
+
+def get_logger(name: str) -> logging.Logger:
+    """Get a logger instance with the given name."""
+    return logging.getLogger(name)

utils/status_utils.py

@@ -0,0 +1,198 @@
+"""
+Status utilities for tracking processing progress and resume capability.
+"""
+
+import os
+from typing import Dict, List, Tuple
+import logging
+
+logger = logging.getLogger(__name__)
+
+def get_processing_status(segments_dir: str) -> Dict[str, any]:
+    """
+    Get detailed processing status for all segments.
+    
+    Args:
+        segments_dir: Directory containing video segments
+        
+    Returns:
+        Dictionary with processing status information
+    """
+    if not os.path.exists(segments_dir):
+        return {
+            'total_segments': 0,
+            'segments_split': 0,
+            'yolo_completed': 0,
+            'sam2_completed': 0,
+            'can_resume': False,
+            'next_step': 'split_video'
+        }
+    
+    # Find all segment directories
+    segments = []
+    for item in os.listdir(segments_dir):
+        item_path = os.path.join(segments_dir, item)
+        if os.path.isdir(item_path) and item.startswith("segment_"):
+            segments.append(item)
+    
+    segments.sort(key=lambda x: int(x.split("_")[1]))
+    
+    # Check status of each segment
+    segments_split = 0
+    yolo_completed = 0
+    sam2_completed = 0
+    
+    for segment in segments:
+        segment_path = os.path.join(segments_dir, segment)
+        segment_idx = int(segment.split("_")[1])
+        
+        # Check if segment video exists
+        video_file = os.path.join(segment_path, f"segment_{str(segment_idx).zfill(3)}.mp4")
+        if os.path.exists(video_file):
+            segments_split += 1
+        
+        # Check if YOLO detection completed
+        yolo_file = os.path.join(segment_path, "yolo_detections")
+        if os.path.exists(yolo_file):
+            yolo_completed += 1
+        
+        # Check if SAM2 processing completed
+        done_file = os.path.join(segment_path, "output_frames_done")
+        if os.path.exists(done_file):
+            sam2_completed += 1
+    
+    # Determine next step
+    next_step = "complete"
+    if sam2_completed < len(segments):
+        next_step = "sam2_processing"
+    elif yolo_completed < len(segments):
+        next_step = "yolo_detection"
+    elif segments_split < len(segments):
+        next_step = "split_video"
+    
+    return {
+        'total_segments': len(segments),
+        'segments_split': segments_split,
+        'yolo_completed': yolo_completed,
+        'sam2_completed': sam2_completed,
+        'can_resume': segments_split > 0,
+        'next_step': next_step,
+        'completion_percentage': (sam2_completed / len(segments) * 100) if segments else 0
+    }
+
+def print_processing_status(segments_dir: str):
+    """
+    Print a human-readable processing status report.
+    
+    Args:
+        segments_dir: Directory containing video segments
+    """
+    status = get_processing_status(segments_dir)
+    
+    print("\\n" + "="*50)
+    print("PROCESSING STATUS REPORT")
+    print("="*50)
+    print(f"Total Segments: {status['total_segments']}")
+    print(f"Video Splitting: {status['segments_split']}/{status['total_segments']} completed")
+    print(f"YOLO Detection: {status['yolo_completed']}/{status['total_segments']} completed")
+    print(f"SAM2 Processing: {status['sam2_completed']}/{status['total_segments']} completed")
+    print(f"Overall Progress: {status['completion_percentage']:.1f}%")
+    print(f"Next Step: {status['next_step']}")
+    print(f"Can Resume: {'Yes' if status['can_resume'] else 'No'}")
+    print("="*50 + "\\n")
+
+def get_incomplete_segments(segments_dir: str) -> List[Tuple[int, str]]:
+    """
+    Get list of segments that still need processing.
+    
+    Args:
+        segments_dir: Directory containing video segments
+        
+    Returns:
+        List of tuples (segment_index, reason)
+    """
+    incomplete = []
+    
+    if not os.path.exists(segments_dir):
+        return incomplete
+    
+    segments = []
+    for item in os.listdir(segments_dir):
+        item_path = os.path.join(segments_dir, item)
+        if os.path.isdir(item_path) and item.startswith("segment_"):
+            segments.append(item)
+    
+    segments.sort(key=lambda x: int(x.split("_")[1]))
+    
+    for segment in segments:
+        segment_path = os.path.join(segments_dir, segment)
+        segment_idx = int(segment.split("_")[1])
+        
+        # Check SAM2 completion first (final step)
+        done_file = os.path.join(segment_path, "output_frames_done")
+        if not os.path.exists(done_file):
+            # Check what step is missing
+            yolo_file = os.path.join(segment_path, "yolo_detections")
+            video_file = os.path.join(segment_path, f"segment_{str(segment_idx).zfill(3)}.mp4")
+            
+            if not os.path.exists(video_file):
+                incomplete.append((segment_idx, "video_splitting"))
+            elif not os.path.exists(yolo_file):
+                incomplete.append((segment_idx, "yolo_detection"))
+            else:
+                incomplete.append((segment_idx, "sam2_processing"))
+    
+    return incomplete
+
+def cleanup_incomplete_segment(segment_dir: str) -> bool:
+    """
+    Clean up a partially processed segment for restart.
+    
+    Args:
+        segment_dir: Path to segment directory
+        
+    Returns:
+        True if cleanup was successful
+    """
+    try:
+        # Remove temporary files that might cause issues
+        temp_files = [
+            "low_res_video.mp4",
+            "output_frames_done"
+        ]
+        
+        removed_count = 0
+        for temp_file in temp_files:
+            temp_path = os.path.join(segment_dir, temp_file)
+            if os.path.exists(temp_path):
+                os.remove(temp_path)
+                removed_count += 1
+        
+        if removed_count > 0:
+            logger.info(f"Cleaned up {removed_count} temporary files from {segment_dir}")
+        
+        return True
+        
+    except Exception as e:
+        logger.error(f"Failed to cleanup segment {segment_dir}: {e}")
+        return False
+
+def find_last_valid_mask(segments_dir: str, before_segment: int) -> str:
+    """
+    Find the most recent segment with a valid mask file.
+    
+    Args:
+        segments_dir: Directory containing segments
+        before_segment: Look for masks before this segment index
+        
+    Returns:
+        Path to the most recent valid mask, or empty string if none found
+    """
+    for i in range(before_segment - 1, -1, -1):
+        segment_path = os.path.join(segments_dir, f"segment_{i}")
+        mask_path = os.path.join(segment_path, "mask.png")
+        
+        if os.path.exists(mask_path):
+            return segment_path
+    
+    return ""