# YOLO + SAM2 Video Processing Pipeline

## Overview

This project provides an automated video processing pipeline that uses YOLO for human detection and SAM2 for precise segmentation to create green screen videos. The system processes long videos by splitting them into manageable segments, detecting and tracking humans in each segment, and then reassembling the processed segments into a final output video with preserved audio.

## Core Functionality

### Input
- **Long video file** (MP4 format, any duration)
- **Configuration file** (YAML format) specifying processing parameters

### Output  
- **Processed video file** with humans visible and background replaced with green screen
- **Preserved audio** from the original input video
- **Intermediate files** for debugging and quality control

## Processing Pipeline

### 1. Video Segmentation
- Splits input video into configurable-duration segments (default: 5 seconds)
- Creates organized directory structure: `segment_0/`, `segment_1/`, etc.
- Each segment folder contains the segment video file
- Generates force keyframes for consistent encoding

### 2. Human Detection & Tracking
- **YOLO Detection**: Automatically detects humans in keyframe segments using YOLOv8
- **SAM2 Segmentation**: Uses detected bounding boxes as prompts for precise mask generation
- **Mask Propagation**: Propagates masks across all frames in each segment
- **Stereo Video Support**: Handles VR/stereo content with left/right human assignment
- **Continuity**: Non-keyframe segments use previous segment masks for consistency

### 3. Green Screen Processing
- **Mask Application**: Applies generated masks to isolate humans
- **Background Replacement**: Replaces non-human areas with green screen (RGB: 0,255,0)
- **GPU Acceleration**: Uses CuPy for fast mask processing
- **Multi-resolution**: Low-res inference for speed, full-res final rendering

### 4. Video Assembly
- **Segment Concatenation**: Combines all processed segments into single video
- **Audio Preservation**: Copies original audio track to final output
- **Quality Maintenance**: Preserves original video quality and framerate

## Key Features

### Automated Processing
- **No Manual Intervention**: Fully automated human detection eliminates manual point selection
- **Batch Processing**: Processes multiple segments efficiently
- **Smart Fallback**: Robust mask propagation with intelligent previous-segment loading

### Modular Architecture
- **Configuration-Driven**: YAML-based configuration for easy parameter adjustment
- **Extensible Design**: Modular structure allows for easy feature additions
- **Error Recovery**: Graceful handling of detection failures and missing segments

### Performance Optimizations
- **GPU Acceleration**: CUDA/NVENC support for faster processing
- **Memory Management**: Efficient handling of large videos through segmentation
- **Concurrent Processing**: Thread-safe operations where applicable

## Technical Stack

### Core Dependencies
- **SAM2**: Facebook's Segment Anything Model 2 for precise segmentation
- **YOLOv8 (Ultralytics)**: Human detection and bounding box generation
- **OpenCV**: Video processing and frame manipulation
- **CuPy**: GPU-accelerated array operations
- **FFmpeg**: Video encoding/decoding and audio handling
- **PyTorch**: Deep learning framework backend

### Supported Formats
- **Input Video**: MP4, AVI, MOV (any OpenCV-supported format)
- **Output Video**: MP4 with H.265/HEVC encoding
- **Audio**: Preserves original audio codec and quality

## Configuration Options

### Video Processing
- `segment_duration`: Duration of each video segment (seconds)
- `inference_scale`: Scale factor for SAM2 inference (for speed)
- `output_scale`: Scale factor for final output

### Detection Parameters
- `yolo_model`: Path to YOLO model weights
- `yolo_confidence`: Detection confidence threshold
- `detect_segments`: Which segments to run YOLO detection on

### SAM2 Parameters
- `sam2_checkpoint`: Path to SAM2 model weights
- `sam2_config`: SAM2 model configuration file

### Output Options
- `use_nvenc`: Enable NVIDIA hardware encoding
- `output_bitrate`: Video bitrate for final output
- `preserve_audio`: Whether to copy audio track

## Directory Structure

```
new_yolo/
├── spec.md                 # This specification document
├── requirements.txt        # Python dependencies
├── config.yaml            # Default configuration file
├── main.py                # Entry point script
├── core/
│   ├── __init__.py
│   ├── video_splitter.py   # Video segmentation logic
│   ├── yolo_detector.py    # YOLO human detection
│   ├── sam2_processor.py   # SAM2 segmentation
│   ├── mask_processor.py   # Mask application and green screen
│   ├── video_assembler.py  # Final video assembly
│   └── config_loader.py    # Configuration management
├── utils/
│   ├── __init__.py
│   ├── file_utils.py       # File system operations
│   ├── video_utils.py      # Video processing utilities
│   └── logging_utils.py    # Logging configuration
└── examples/
    ├── basic_config.yaml   # Example configuration
    └── advanced_config.yaml # Advanced configuration options
```

## Usage Examples

### Basic Usage
```bash
python main.py --config config.yaml
```

### Custom Configuration
```bash
python main.py --config examples/advanced_config.yaml
```

### Configuration File Example
```yaml
input:
  video_path: "/path/to/input/video.mp4"
  
output:
  directory: "/path/to/output/"
  filename: "processed_video.mp4"
  
processing:
  segment_duration: 5
  inference_scale: 0.5
  yolo_confidence: 0.6
  detect_segments: "all"  # or [0, 5, 10]
  
models:
  yolo_model: "yolov8n.pt"
  sam2_checkpoint: "../checkpoints/sam2.1_hiera_large.pt"
  sam2_config: "configs/sam2.1/sam2.1_hiera_l.yaml"
```

## Use Cases

### Content Creation
- **VR/360 Video Processing**: Remove backgrounds from immersive content
- **Green Screen Production**: Automated background removal for video production
- **Social Media Content**: Quick background replacement for content creators

### Commercial Applications  
- **Video Conferencing**: Real-time background replacement
- **E-learning**: Professional video production with clean backgrounds
- **Marketing**: Product demonstration videos with custom backgrounds

## Performance Considerations

### Hardware Requirements
- **GPU**: NVIDIA GPU with CUDA support (recommended)
- **RAM**: 16GB+ for processing large videos
- **Storage**: SSD recommended for temporary file operations

### Processing Time
- Approximately **1-2x real-time** on modern GPUs
- Scales with video resolution and segment count
- Memory usage remains constant regardless of input video length

## Future Enhancements

### Planned Features
- **Multi-object Tracking**: Support for multiple humans per frame
- **Custom Object Detection**: Configurable object classes beyond humans
- **Real-time Processing**: Live video stream support
- **Cloud Integration**: AWS/GCP processing support
- **Web Interface**: Browser-based configuration and monitoring

### Model Improvements
- **Fine-tuned YOLO**: Domain-specific human detection models
- **SAM2 Optimization**: Custom SAM2 checkpoints for video content
- **Temporal Consistency**: Enhanced cross-segment mask propagation


Here is the original monolithic script this repo is a refactor/modularization of. If something
doesn't work in this repo, then consult the following script becasue it does work so this can 
be used to solve problems:


import os
import cv2
import numpy as np
import cupy as cp
from concurrent.futures import ThreadPoolExecutor
import torch
import logging
import sys
import gc
from sam2.build_sam import build_sam2_video_predictor
import argparse
from ultralytics import YOLO

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

# Variables for input and output directories
SAM2_CHECKPOINT = "../checkpoints/sam2.1_hiera_large.pt"
MODEL_CFG = "configs/sam2.1/sam2.1_hiera_l.yaml"
GREEN = [0, 255, 0]
BLUE = [255, 0, 0]

INFERENCE_SCALE = 0.50
FULL_SCALE = 1.0

# YOLO model for human detection (class 0 = person)
YOLO_MODEL_PATH = "yolov8n.pt"  # You can change this to a custom model
YOLO_CONFIDENCE = 0.6
HUMAN_CLASS_ID = 0  # COCO class ID for person

def open_video(video_path):
    """
    Opens a video file and returns a generator that yields frames.

    Parameters:
    - video_path: Path to the video file.

    Returns:
    - A generator that yields frames from the video.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Could not open video file {video_path}")
        return
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        yield frame
    cap.release()

def load_previous_segment_mask(prev_segment_dir):
    mask_path = os.path.join(prev_segment_dir,  "mask.png")
    mask_image = cv2.imread(mask_path)

    if mask_image is None:
        raise FileNotFoundError(f"Mask image not found at {mask_path}")

    # Ensure the mask_image has three color channels
    if len(mask_image.shape) != 3 or mask_image.shape[2] != 3:
        raise ValueError("Mask image does not have three color channels.")

    mask_image = mask_image.astype(np.uint8)
    
    # Extract Object A and Object B masks
    mask_a = np.all(mask_image == GREEN, axis=2)
    mask_b = np.all(mask_image == BLUE, axis=2)

    per_obj_input_mask = {1: mask_a, 2: mask_b}
    input_palette = None  # No palette needed for binary mask
    
    return per_obj_input_mask, input_palette


def apply_green_mask(frame, masks):
    # Convert frame and masks to CuPy arrays
    frame_gpu = cp.asarray(frame)
    combined_mask = cp.zeros(frame_gpu.shape[:2], dtype=cp.bool_)

    for mask in masks:
        mask_gpu = cp.asarray(mask.squeeze())
        if mask_gpu.shape != frame_gpu.shape[:2]:
            resized_mask = cv2.resize(cp.asnumpy(mask_gpu).astype(cp.float32), 
                                      (frame_gpu.shape[1], frame_gpu.shape[0]))
            mask_gpu = cp.asarray(resized_mask > 0.5)  # Convert back to CuPy boolean array
        else:
            mask_gpu = mask_gpu.astype(cp.bool_)  # Ensure boolean type
        combined_mask |= mask_gpu  # Perform the bitwise OR operation

    green_background = cp.full(frame_gpu.shape, cp.array([0, 255, 0], dtype=cp.uint8), dtype=cp.uint8)
    result_frame = cp.where(combined_mask[..., None], frame_gpu, green_background)
    return cp.asnumpy(result_frame)  # Convert back to NumPy


def initialize_predictor():
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print(
            "\nSupport for MPS devices is preliminary. SAM 2 is trained with CUDA and might "
            "give numerically different outputs and sometimes degraded performance on MPS."
        )
        # Enable MPS fallback for operations not supported on MPS
        os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
    else:
        device = torch.device("cpu")
    logger.info(f"Using device: {device}")
    predictor = build_sam2_video_predictor(MODEL_CFG, SAM2_CHECKPOINT, device=device)
    return predictor


def load_first_frame(video_path, scale=1.0):
    """
    Opens a video file and returns the first frame, scaled as specified.

    Parameters:
    - video_path: Path to the video file.
    - scale: Scaling factor for the frame (default is 1.0 for original size).

    Returns:
    - first_frame: The first frame of the video, scaled accordingly.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        logger.error(f"Error: Could not open video file {video_path}")
        return None

    ret, frame = cap.read()
    cap.release()

    if not ret:
        logger.error(f"Error: Could not read frame from video file {video_path}")
        return None

    if scale != 1.0:
        frame = cv2.resize(
            frame, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR
        )

    return frame

def detect_humans_with_yolo(frame, yolo_model, confidence_threshold=YOLO_CONFIDENCE):
    """
    Detect humans in a frame using YOLO model.
    
    Parameters:
    - frame: Input frame (BGR format)
    - yolo_model: Loaded YOLO model
    - confidence_threshold: Detection confidence threshold
    
    Returns:
    - human_boxes: List of bounding boxes for detected humans
    """
    # Run YOLO detection
    results = yolo_model(frame, conf=confidence_threshold, verbose=False)
    
    human_boxes = []
    
    # Process results
    for result in results:
        boxes = result.boxes
        if boxes is not None:
            for box in boxes:
                # Get class ID
                cls = int(box.cls.cpu().numpy()[0])
                
                # Check if it's a person (class 0 in COCO)
                if cls == HUMAN_CLASS_ID:
                    # Get bounding box coordinates (x1, y1, x2, y2)
                    coords = box.xyxy[0].cpu().numpy()
                    conf = float(box.conf.cpu().numpy()[0])
                    
                    human_boxes.append({
                        'bbox': coords,
                        'confidence': conf
                    })
                    
                    logger.info(f"Detected human with confidence {conf:.2f} at {coords}")
    
    return human_boxes

def add_yolo_detections_to_predictor(predictor, inference_state, human_detections, frame_width):
    """
    Add YOLO human detections as bounding boxes to SAM2 predictor.
    For stereo videos, creates two objects (left and right humans).
    
    Parameters:
    - predictor: SAM2 video predictor
    - inference_state: SAM2 inference state
    - human_detections: List of human detection results
    - frame_width: Width of the frame for stereo splitting
    
    Returns:
    - out_mask_logits: SAM2 output mask logits
    """
    half_frame_width = frame_width // 2
    
    # Sort detections by x-coordinate to get left and right humans
    human_detections.sort(key=lambda x: x['bbox'][0])  # Sort by x1 coordinate
    
    obj_id = 1
    out_mask_logits = None
    
    for i, detection in enumerate(human_detections[:2]):  # Take up to 2 humans (left and right)
        bbox = detection['bbox']
        
        # For stereo videos, assign obj_id based on position
        if len(human_detections) >= 2:
            # If we have multiple humans, assign based on left/right position
            center_x = (bbox[0] + bbox[2]) / 2
            if center_x < half_frame_width:
                current_obj_id = 1  # Left human
            else:
                current_obj_id = 2  # Right human
        else:
            # If only one human, duplicate for both sides (as in original stereo logic)
            current_obj_id = obj_id
            obj_id += 1
            
            # Also add the mirrored version for stereo
            if obj_id <= 2:
                mirrored_bbox = bbox.copy()
                mirrored_bbox[0] += half_frame_width  # Shift x1
                mirrored_bbox[2] += half_frame_width  # Shift x2
                
                # Ensure mirrored bbox is within frame bounds
                mirrored_bbox[0] = max(0, min(mirrored_bbox[0], frame_width - 1))
                mirrored_bbox[2] = max(0, min(mirrored_bbox[2], frame_width - 1))
                
                try:
                    _, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box(
                        inference_state=inference_state,
                        frame_idx=0,
                        obj_id=obj_id,
                        box=mirrored_bbox.astype(np.float32),
                    )
                    logger.info(f"Added mirrored human detection for Object {obj_id}")
                    obj_id += 1
                except Exception as e:
                    logger.error(f"Error adding mirrored human detection for Object {obj_id}: {e}")
        
        try:
            _, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box(
                inference_state=inference_state,
                frame_idx=0,
                obj_id=current_obj_id,
                box=bbox.astype(np.float32),
            )
            logger.info(f"Added human detection for Object {current_obj_id}")
        except Exception as e:
            logger.error(f"Error adding human detection for Object {current_obj_id}: {e}")
    
    return out_mask_logits

def propagate_masks(predictor, inference_state):
    video_segments = {}
    for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
        video_segments[out_frame_idx] = {
            out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
            for i, out_obj_id in enumerate(out_obj_ids)
        }
    return video_segments

def apply_colored_mask(frame, masks_a, masks_b):
    colored_mask = np.zeros_like(frame)
    
    # Apply colors to the masks
    for mask in masks_a:
        mask = mask.squeeze()
        if mask.shape != frame.shape[:2]:
            mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
        indices = np.where(mask)
        colored_mask[mask] = [0, 255, 0]  # Green for Object A
    
    for mask in masks_b:
        mask = mask.squeeze()
        if mask.shape != frame.shape[:2]:
            mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
        indices = np.where(mask)
        colored_mask[mask] = [255, 0, 0]  # Blue for Object B
    
    return colored_mask


def process_and_save_output_video(video_path, output_video_path, video_segments, use_nvenc=False):
    """
    Process high-resolution frames, apply upscaled masks, and save the output video.
    """
    cap = cv2.VideoCapture(video_path)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 59.94

    # Setup VideoWriter with desired settings
    if use_nvenc:
        # Use FFmpeg with NVENC offloading for H.265 encoding
        import subprocess

        if sys.platform == 'darwin':
            encoder = 'hevc_videotoolbox'
        else:
            encoder = 'hevc_nvenc'

        command = [
            'ffmpeg',
            '-y',  # Overwrite output file if it exists
            '-f', 'rawvideo',
            '-vcodec', 'rawvideo',
            '-pix_fmt', 'bgr24',
            '-s', f'{frame_width}x{frame_height}',
            '-r', str(fps),
            '-i', '-',  # Input from stdin
            '-an',  # No audio
            '-vcodec', encoder,
            '-pix_fmt', 'nv12',
            '-preset', 'slow',
            '-b:v', '50M',
            output_video_path
        ]
        process = subprocess.Popen(command, stdin=subprocess.PIPE)
    else:
        # Use OpenCV VideoWriter
        fourcc = cv2.VideoWriter_fourcc(*'HEVC')  # H.265
        out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret or frame_idx >= len(video_segments):
            break

        masks = [video_segments[frame_idx][out_obj_id] for out_obj_id in video_segments[frame_idx]]
        upscaled_masks = []

        for mask in masks:
            mask = mask.squeeze()
            upscaled_mask = cv2.resize(mask.astype(np.uint8), (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
            upscaled_masks.append(upscaled_mask)

        result_frame = apply_green_mask(frame, upscaled_masks)

        # Write frame to output
        if use_nvenc:
            process.stdin.write(result_frame.tobytes())
        else:
            out.write(result_frame)

        frame_idx += 1

    cap.release()
    if use_nvenc:
        process.stdin.close()
        process.wait()
    else:
        out.release()

def get_video_file_name(index):
    return f"segment_{str(index).zfill(3)}.mp4"

def do_yolo_detection_on_segments(base_dir, segments, detect_segments, scale=1.0, yolo_model_path=YOLO_MODEL_PATH):
    """
    Run YOLO detection on specified segments and save detection results.
    """
    logger.info("Running YOLO detection on requested segments.")
    
    # Load YOLO model
    yolo_model = YOLO(yolo_model_path)
    
    for i, segment in enumerate(segments):
        segment_index = int(segment.split("_")[1])
        segment_dir = os.path.join(base_dir, segment)
        detection_file = os.path.join(segment_dir, "yolo_detections")
        video_file = os.path.join(segment_dir, get_video_file_name(i))
        
        if segment_index in detect_segments and not os.path.exists(detection_file):
            first_frame = load_first_frame(video_file, scale)
            if first_frame is None:
                continue
                
            # Convert BGR to RGB for YOLO (YOLO expects BGR, so keep as BGR)
            human_detections = detect_humans_with_yolo(first_frame, yolo_model)
            
            if human_detections:
                # Save detection results
                with open(detection_file, 'w') as f:
                    f.write("# YOLO Human Detections\n")
                    for detection in human_detections:
                        bbox = detection['bbox']
                        conf = detection['confidence']
                        f.write(f"{bbox[0]},{bbox[1]},{bbox[2]},{bbox[3]},{conf}\n")
                logger.info(f"Saved {len(human_detections)} human detections for segment {segment}")
            else:
                logger.warning(f"No humans detected in segment {segment}")
                # Create empty file to mark as processed
                with open(detection_file, 'w') as f:
                    f.write("# No humans detected\n")

def save_final_masks(video_segments, mask_output_path):
    """
    Save the final masks as a colored image.
    """
    last_frame_idx = max(video_segments.keys())
    masks_dict = video_segments[last_frame_idx]
    # Assuming you have two objects with IDs 1 and 2
    mask_a = masks_dict.get(1).squeeze() if 1 in masks_dict else None
    mask_b = masks_dict.get(2).squeeze() if 2 in masks_dict else None
    
    if mask_a is None and mask_b is None:
        logger.error("No masks found for objects.")
        return
        
    # Use the first available mask to determine dimensions
    reference_mask = mask_a if mask_a is not None else mask_b
    black_frame = np.zeros((reference_mask.shape[0], reference_mask.shape[1], 3), dtype=np.uint8)
    
    if mask_a is not None:
        mask_a = mask_a.astype(bool)
        black_frame[mask_a] = GREEN
        
    if mask_b is not None:
        mask_b = mask_b.astype(bool) 
        black_frame[mask_b] = BLUE

    # Save the mask image
    cv2.imwrite(mask_output_path, black_frame)
    logger.info(f"Saved final masks to {mask_output_path}")

def create_low_res_video(input_video_path, output_video_path, scale):
    """
    Creates a low-resolution version of the input video for inference.
    """
    cap = cv2.VideoCapture(input_video_path)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) * scale)
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * scale)
    fps = cap.get(cv2.CAP_PROP_FPS) or 59.94

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        low_res_frame = cv2.resize(frame, (frame_width, frame_height), interpolation=cv2.INTER_LINEAR)
        out.write(low_res_frame)

    cap.release()
    out.release()


def main():
    parser = argparse.ArgumentParser(description="Process video segments with YOLO + SAM2.")
    parser.add_argument("--base-dir", type=str, help="Base directory for video segments.")
    parser.add_argument("--segments-detect-humans", nargs='*', help="Segments for which to run YOLO human detection. Use 'all' for all segments, or list specific segment numbers (e.g., 1 5 10). Default: all segments.")
    parser.add_argument("--yolo-model", type=str, default=YOLO_MODEL_PATH, help="Path to YOLO model.")
    parser.add_argument("--yolo-confidence", type=float, default=YOLO_CONFIDENCE, help="YOLO detection confidence threshold.")
    args = parser.parse_args()

    base_dir = args.base_dir
    segments = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d)) and d.startswith("segment_")]
    segments.sort(key=lambda x: int(x.split("_")[1]))
    
    # Handle different ways to specify segments for YOLO detection
    if args.segments_detect_humans is None or len(args.segments_detect_humans) == 0:
        # Default: run YOLO on all segments
        detect_segments = [int(seg.split("_")[1]) for seg in segments]
        logger.info("No segments specified, running YOLO detection on ALL segments")
    elif len(args.segments_detect_humans) == 1 and args.segments_detect_humans[0].lower() == 'all':
        # Explicit 'all' keyword
        detect_segments = [int(seg.split("_")[1]) for seg in segments]
        logger.info("Running YOLO detection on ALL segments")
    else:
        # Specific segment numbers provided
        try:
            detect_segments = [int(x) for x in args.segments_detect_humans]
            logger.info(f"Running YOLO detection on segments: {detect_segments}")
        except ValueError:
            logger.error("Invalid segment numbers provided. Use integers or 'all'.")
            return
    
    # Run YOLO detection on specified segments
    do_yolo_detection_on_segments(base_dir, segments, detect_segments, scale=INFERENCE_SCALE, yolo_model_path=args.yolo_model)
    
    # Load YOLO model for inference
    yolo_model = YOLO(args.yolo_model)

    for i, segment in enumerate(segments):
        segment_index = int(segment.split("_")[1])
        segment_dir = os.path.join(base_dir, segment)
        video_file_name = get_video_file_name(i)
        video_path = os.path.join(segment_dir, video_file_name)
        output_done_file = os.path.join(segment_dir, "output_frames_done")
        
        if os.path.exists(output_done_file):
            logger.info(f"Segment {segment} already processed. Skipping.")
            continue

        logger.info(f"Processing segment {segment}")

        # Initialize predictor
        predictor = initialize_predictor()

        # Prepare low-resolution video frames for inference
        low_res_video_path = os.path.join(segment_dir, "low_res_video.mp4")
        if not os.path.exists(low_res_video_path):
            create_low_res_video(video_path, low_res_video_path, INFERENCE_SCALE)
            logger.info(f"Low-resolution video created for segment {segment}")
        else:
            logger.info(f"Low-resolution video already exists for segment {segment}, reuse")

        # Initialize inference state with low-resolution video
        inference_state = predictor.init_state(video_path=low_res_video_path, async_loading_frames=True)

        # Load YOLO detections or previous masks
        detection_file = os.path.join(segment_dir, "yolo_detections")
        use_detections = segment_index in detect_segments

        if i == 0 and not use_detections:
            # First segment must use YOLO detection since there's no previous mask
            logger.warning(f"First segment {segment} requires YOLO detection. Running YOLO detection.")
            use_detections = True
            
        if i > 0 and not use_detections:
            # Try to load previous segment mask - search backwards for the most recent successful mask
            logger.info(f"Using previous segment mask for segment {segment}")
            mask_found = False
            
            # Search backwards through previous segments to find a valid mask
            for j in range(i - 1, -1, -1):
                prev_segment_dir = os.path.join(base_dir, segments[j])
                prev_mask_path = os.path.join(prev_segment_dir, "mask.png")
                
                if os.path.exists(prev_mask_path):
                    try:
                        per_obj_input_mask, input_palette = load_previous_segment_mask(prev_segment_dir)
                        # Add previous masks to predictor
                        for obj_id, mask in per_obj_input_mask.items():
                            predictor.add_new_mask(inference_state, 0, obj_id, mask)
                        logger.info(f"Successfully loaded mask from segment {segments[j]}")
                        mask_found = True
                        break
                    except Exception as e:
                        logger.warning(f"Error loading mask from {segments[j]}: {e}")
                        continue
            
            if not mask_found:
                logger.error(f"No valid previous mask found for segment {segment}. Consider running YOLO detection on this segment.")
                continue
        else:
            # Load first frame for detection
            first_frame = load_first_frame(low_res_video_path, scale=1.0)
            if first_frame is None:
                logger.error(f"Could not load first frame for segment {segment}")
                continue
                
            # Run YOLO detection on first frame (either from file or on-the-fly)
            if os.path.exists(detection_file):
                logger.info(f"Using existing YOLO detections for segment {segment}")
            else:
                logger.info(f"Running YOLO detection on-the-fly for segment {segment}")
                
            human_detections = detect_humans_with_yolo(first_frame, yolo_model, args.yolo_confidence)
            
            if human_detections:
                # Add YOLO detections to predictor
                frame_width = first_frame.shape[1]
                add_yolo_detections_to_predictor(predictor, inference_state, human_detections, frame_width)
            else:
                logger.warning(f"No humans detected in segment {segment}")
                continue

        # Perform inference and collect masks per frame
        video_segments = propagate_masks(predictor, inference_state)

        # Process high-resolution frames and save output video
        output_video_path = os.path.join(segment_dir, f"output_{segment_index}.mp4")
        logger.info("Processing segment complete, attempting to save full video from low res masks")
        process_and_save_output_video(
            video_path,
            output_video_path,
            video_segments,
            use_nvenc=True  # Set to True to use NVENC offloading
        )

        # Save final masks
        mask_output_path = os.path.join(segment_dir, "mask.png")
        save_final_masks(video_segments, mask_output_path)

        # Clean up
        predictor.reset_state(inference_state)
        del inference_state
        del video_segments
        del predictor
        gc.collect()

        try:
            os.remove(low_res_video_path)
            logger.info(f"Deleted low-resolution video for segment {segment}")
        except Exception as e:
            logger.warning(f"Could not delete low-resolution video for segment {segment}: {e}")

        # Mark segment as completed
        open(output_done_file, 'a').close()
        
    logger.info("Processing complete.")

if __name__ == "__main__":
    main()