main.py

#!/usr/bin/env python3
import cv2, os, subprocess, argparse
from PIL import Image
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, SamModel, SamProcessor
from tqdm import tqdm
import numpy as np
from datetime import datetime
from deep_sort_integration import DeepSORTTracker
from scenedetect import detect, ContentDetector
from functools import lru_cache

# Constants
DEFAULT_TEST_MODE_DURATION = 3  # Process only first 3 seconds in test mode by default
FFMPEG_PRESETS = [
    "ultrafast",
    "superfast",
    "veryfast",
    "faster",
    "fast",
    "medium",
    "slow",
    "slower",
    "veryslow",
]
FONT = cv2.FONT_HERSHEY_SIMPLEX  # Font for bounding-box-style labels

# Detection parameters
IOU_THRESHOLD = 0.5  # IoU threshold for considering boxes related

# Hitmarker parameters
HITMARKER_SIZE = 20  # Size of the hitmarker in pixels
HITMARKER_GAP = 3  # Size of the empty space in the middle (reduced from 8)
HITMARKER_THICKNESS = 2  # Thickness of hitmarker lines
HITMARKER_COLOR = (255, 255, 255)  # White color for hitmarker
HITMARKER_SHADOW_COLOR = (80, 80, 80)  # Lighter gray for shadow effect
HITMARKER_SHADOW_OFFSET = 1  # Smaller shadow offset

# SAM parameters
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Initialize model variables as None
sam_model = None
sam_processor = None
slimsam_model = None
slimsam_processor = None

@lru_cache(maxsize=2)  # Cache both regular and slim SAM models
def get_sam_model(slim=False):
    """Get cached SAM model and processor."""
    global sam_model, sam_processor, slimsam_model, slimsam_processor
    
    if slim:
        if slimsam_model is None:
            print("Loading SlimSAM model for the first time...")
            slimsam_model = SamModel.from_pretrained("nielsr/slimsam-50-uniform").to(device)
            slimsam_processor = SamProcessor.from_pretrained("nielsr/slimsam-50-uniform")
        return slimsam_model, slimsam_processor
    else:
        if sam_model is None:
            print("Loading SAM model for the first time...")
            sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
            sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
        return sam_model, sam_processor

def load_sam_model(slim=False):
    """Load SAM model and processor with caching."""
    return get_sam_model(slim=slim)

def generate_color_pair():
    """Generate a generic light blue and dark blue color pair for SAM visualization."""
    dark_rgb = [0, 0, 139]  # Dark blue
    light_rgb = [173, 216, 230]  # Light blue
    return dark_rgb, light_rgb

def create_mask_overlay(image, masks, points=None, labels=None):
    """Create a mask overlay with contours for multiple SAM visualizations.
    
    Args:
        image: PIL Image to overlay masks on
        masks: List of binary masks or single mask
        points: Optional list of (x,y) points for labels
        labels: Optional list of label strings for each point
    """
    # Convert single mask to list for uniform processing
    if not isinstance(masks, list):
        masks = [masks]
    
    # Create empty overlays
    overlay = np.zeros((*image.size[::-1], 4), dtype=np.uint8)
    outline = np.zeros((*image.size[::-1], 4), dtype=np.uint8)
    
    # Process each mask
    for i, mask in enumerate(masks):
        # Convert binary mask to uint8
        mask_uint8 = (mask > 0).astype(np.uint8)
        
        # Dilation to fill gaps
        kernel = np.ones((5, 5), np.uint8)
        mask_dilated = cv2.dilate(mask_uint8, kernel, iterations=1)
        
        # Find contours of the dilated mask
        contours, _ = cv2.findContours(mask_dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # Generate random color pair for this segmentation
        dark_color, light_color = generate_color_pair()
        
        # Add to the overlays
        overlay[mask_dilated > 0] = [*light_color, 90]  # Light color with 35% opacity
        cv2.drawContours(outline, contours, -1, (*dark_color, 255), 2)  # Dark color outline
    
    # Convert to PIL images
    mask_overlay = Image.fromarray(overlay, 'RGBA')
    outline_overlay = Image.fromarray(outline, 'RGBA')
    
    # Composite the layers
    result = image.convert('RGBA')
    result.paste(mask_overlay, (0, 0), mask_overlay)
    result.paste(outline_overlay, (0, 0), outline_overlay)
    
    # Add labels if provided
    if points and labels:
        result_array = np.array(result)
        for (x, y), label in zip(points, labels):
            label_size = cv2.getTextSize(label, FONT, 0.5, 1)[0]
            cv2.putText(
                result_array,
                label,
                (int(x - label_size[0] // 2), int(y - 20)),
                FONT,
                0.5,
                (255, 255, 255),
                1,
                cv2.LINE_AA,
            )
        result = Image.fromarray(result_array)
    
    return result

def process_sam_detection(image, center_x, center_y, slim=False):
    """Process a single detection point with SAM.
    
    Returns:
        tuple: (mask, result_pil) where mask is the binary mask and result_pil is the visualization
    """
    if not isinstance(image, Image.Image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    
    # Get appropriate model from cache
    model, processor = get_sam_model(slim)
    
    # Process the image with SAM
    inputs = processor(
        image,
        input_points=[[[center_x, center_y]]],
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        outputs = model(**inputs)

    mask = processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"].cpu(),
        inputs["reshaped_input_sizes"].cpu()
    )[0][0][0].numpy()
    
    # Create the visualization
    result = create_mask_overlay(image, mask)
    return mask, result

def load_moondream():
    """Load Moondream model and tokenizer."""
    model = AutoModelForCausalLM.from_pretrained(
        "vikhyatk/moondream2", trust_remote_code=True, device_map={"": "cuda"}
    )
    tokenizer = AutoTokenizer.from_pretrained("vikhyatk/moondream2")
    return model, tokenizer


def get_video_properties(video_path):
    """Get basic video properties."""
    video = cv2.VideoCapture(video_path)
    fps = video.get(cv2.CAP_PROP_FPS)
    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    video.release()
    return {"fps": fps, "frame_count": frame_count, "width": width, "height": height}


def is_valid_bounding_box(bounding_box):
    """Check if bounding box coordinates are reasonable."""
    x1, y1, x2, y2 = bounding_box
    width = x2 - x1
    height = y2 - y1

    # Reject boxes that are too large (over 90% of frame in both dimensions)
    if width > 0.9 and height > 0.9:
        return False

    # Reject boxes that are too small (less than 1% of frame)
    if width < 0.01 or height < 0.01:
        return False

    return True


def split_frame_into_grid(frame, grid_rows, grid_cols):
    """Split a frame into a grid of tiles."""
    height, width = frame.shape[:2]
    tile_height = height // grid_rows
    tile_width = width // grid_cols
    tiles = []
    tile_positions = []

    for i in range(grid_rows):
        for j in range(grid_cols):
            y1 = i * tile_height
            y2 = (i + 1) * tile_height if i < grid_rows - 1 else height
            x1 = j * tile_width
            x2 = (j + 1) * tile_width if j < grid_cols - 1 else width

            tile = frame[y1:y2, x1:x2]
            tiles.append(tile)
            tile_positions.append((x1, y1, x2, y2))

    return tiles, tile_positions


def convert_tile_coords_to_frame(box, tile_pos, frame_shape):
    """Convert coordinates from tile space to frame space."""
    frame_height, frame_width = frame_shape[:2]
    tile_x1, tile_y1, tile_x2, tile_y2 = tile_pos
    tile_width = tile_x2 - tile_x1
    tile_height = tile_y2 - tile_y1

    x1_tile_abs = box[0] * tile_width
    y1_tile_abs = box[1] * tile_height
    x2_tile_abs = box[2] * tile_width
    y2_tile_abs = box[3] * tile_height

    x1_frame_abs = tile_x1 + x1_tile_abs
    y1_frame_abs = tile_y1 + y1_tile_abs
    x2_frame_abs = tile_x1 + x2_tile_abs
    y2_frame_abs = tile_y1 + y2_tile_abs

    x1_norm = x1_frame_abs / frame_width
    y1_norm = y1_frame_abs / frame_height
    x2_norm = x2_frame_abs / frame_width
    y2_norm = y2_frame_abs / frame_height

    x1_norm = max(0.0, min(1.0, x1_norm))
    y1_norm = max(0.0, min(1.0, y1_norm))
    x2_norm = max(0.0, min(1.0, x2_norm))
    y2_norm = max(0.0, min(1.0, y2_norm))

    return [x1_norm, y1_norm, x2_norm, y2_norm]


def merge_tile_detections(tile_detections, iou_threshold=0.5):
    """Merge detections from different tiles using NMS-like approach."""
    if not tile_detections:
        return []

    all_boxes = []
    all_keywords = []

    # Collect all boxes and their keywords
    for detections in tile_detections:
        for box, keyword in detections:
            all_boxes.append(box)
            all_keywords.append(keyword)

    if not all_boxes:
        return []

    # Convert to numpy for easier processing
    boxes = np.array(all_boxes)

    # Calculate areas
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)

    # Sort boxes by area
    order = areas.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        if order.size == 1:
            break

        # Calculate IoU with rest of boxes
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)
        inter = w * h

        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        # Get indices of boxes with IoU less than threshold
        inds = np.where(ovr <= iou_threshold)[0]
        order = order[inds + 1]

    return [(all_boxes[i], all_keywords[i]) for i in keep]


def detect_objects_in_frame(model, tokenizer, image, target_object, grid_rows=1, grid_cols=1):
    """Detect specified objects in a frame using grid-based analysis."""
    if grid_rows == 1 and grid_cols == 1:
        return detect_objects_in_frame_single(model, tokenizer, image, target_object)

    # Convert numpy array to PIL Image if needed
    if not isinstance(image, Image.Image):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Split frame into tiles
    tiles, tile_positions = split_frame_into_grid(image, grid_rows, grid_cols)

    # Process each tile
    tile_detections = []
    for tile, tile_pos in zip(tiles, tile_positions):
        # Convert tile to PIL Image
        tile_pil = Image.fromarray(tile)

        # Detect objects in tile
        response = model.detect(tile_pil, target_object)

        if response and "objects" in response and response["objects"]:
            objects = response["objects"]
            tile_objects = []

            for obj in objects:
                if all(k in obj for k in ["x_min", "y_min", "x_max", "y_max"]):
                    box = [obj["x_min"], obj["y_min"], obj["x_max"], obj["y_max"]]

                    if is_valid_bounding_box(box):
                        # Convert tile coordinates to frame coordinates
                        frame_box = convert_tile_coords_to_frame(
                            box, tile_pos, image.shape
                        )
                        tile_objects.append((frame_box, target_object))

            if tile_objects:  # Only append if we found valid objects
                tile_detections.append(tile_objects)

    # Merge detections from all tiles
    merged_detections = merge_tile_detections(tile_detections)
    return merged_detections


def detect_objects_in_frame_single(model, tokenizer, image, target_object):
    """Single-frame detection function."""
    detected_objects = []

    # Convert numpy array to PIL Image if needed
    if not isinstance(image, Image.Image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    # Detect objects
    response = model.detect(image, target_object)

    # Check if we have valid objects
    if response and "objects" in response and response["objects"]:
        objects = response["objects"]

        for obj in objects:
            if all(k in obj for k in ["x_min", "y_min", "x_max", "y_max"]):
                box = [obj["x_min"], obj["y_min"], obj["x_max"], obj["y_max"]]
                # If box is valid (not full-frame), add it
                if is_valid_bounding_box(box):
                    detected_objects.append((box, target_object))

    return detected_objects


def draw_hitmarker(
    frame, center_x, center_y, size=HITMARKER_SIZE, color=HITMARKER_COLOR, shadow=True
):
    """Draw a COD-style hitmarker cross with more space in the middle."""
    half_size = size // 2

    # Draw shadow first if enabled
    if shadow:
        # Top-left to center shadow
        cv2.line(
            frame,
            (
                center_x - half_size + HITMARKER_SHADOW_OFFSET,
                center_y - half_size + HITMARKER_SHADOW_OFFSET,
            ),
            (
                center_x - HITMARKER_GAP + HITMARKER_SHADOW_OFFSET,
                center_y - HITMARKER_GAP + HITMARKER_SHADOW_OFFSET,
            ),
            HITMARKER_SHADOW_COLOR,
            HITMARKER_THICKNESS,
        )
        # Top-right to center shadow
        cv2.line(
            frame,
            (
                center_x + half_size + HITMARKER_SHADOW_OFFSET,
                center_y - half_size + HITMARKER_SHADOW_OFFSET,
            ),
            (
                center_x + HITMARKER_GAP + HITMARKER_SHADOW_OFFSET,
                center_y - HITMARKER_GAP + HITMARKER_SHADOW_OFFSET,
            ),
            HITMARKER_SHADOW_COLOR,
            HITMARKER_THICKNESS,
        )
        # Bottom-left to center shadow
        cv2.line(
            frame,
            (
                center_x - half_size + HITMARKER_SHADOW_OFFSET,
                center_y + half_size + HITMARKER_SHADOW_OFFSET,
            ),
            (
                center_x - HITMARKER_GAP + HITMARKER_SHADOW_OFFSET,
                center_y + HITMARKER_GAP + HITMARKER_SHADOW_OFFSET,
            ),
            HITMARKER_SHADOW_COLOR,
            HITMARKER_THICKNESS,
        )
        # Bottom-right to center shadow
        cv2.line(
            frame,
            (
                center_x + half_size + HITMARKER_SHADOW_OFFSET,
                center_y + half_size + HITMARKER_SHADOW_OFFSET,
            ),
            (
                center_x + HITMARKER_GAP + HITMARKER_SHADOW_OFFSET,
                center_y + HITMARKER_GAP + HITMARKER_SHADOW_OFFSET,
            ),
            HITMARKER_SHADOW_COLOR,
            HITMARKER_THICKNESS,
        )

    # Draw main hitmarker
    # Top-left to center
    cv2.line(
        frame,
        (center_x - half_size, center_y - half_size),
        (center_x - HITMARKER_GAP, center_y - HITMARKER_GAP),
        color,
        HITMARKER_THICKNESS,
    )
    # Top-right to center
    cv2.line(
        frame,
        (center_x + half_size, center_y - half_size),
        (center_x + HITMARKER_GAP, center_y - HITMARKER_GAP),
        color,
        HITMARKER_THICKNESS,
    )
    # Bottom-left to center
    cv2.line(
        frame,
        (center_x - half_size, center_y + half_size),
        (center_x - HITMARKER_GAP, center_y + HITMARKER_GAP),
        color,
        HITMARKER_THICKNESS,
    )
    # Bottom-right to center
    cv2.line(
        frame,
        (center_x + half_size, center_y + half_size),
        (center_x + HITMARKER_GAP, center_y + HITMARKER_GAP),
        color,
        HITMARKER_THICKNESS,
    )


def draw_ad_boxes(frame, detected_objects, detect_keyword, model, box_style="censor"):
    height, width = frame.shape[:2]

    points = []
    # Only get points if we need them for hitmarker or SAM styles
    if box_style in ["hitmarker", "sam", "sam-fast"]:
        frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        try:
            point_response = model.point(frame_pil, detect_keyword)
            
            if isinstance(point_response, dict) and 'points' in point_response:
                points = point_response['points']
        except Exception as e:
            print(f"Error during point detection: {str(e)}")
            points = []

    # Only load SAM models and process points if we're using SAM styles and have points
    if box_style in ["sam", "sam-fast"] and points:
        # Start with the original PIL image
        frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        
        # Collect all masks and points
        all_masks = []
        point_coords = []
        point_labels = []
        
        for point in points:
            try:
                center_x = int(float(point["x"]) * width)
                center_y = int(float(point["y"]) * height)

                # Get mask and visualization
                mask, _ = process_sam_detection(frame_pil, center_x, center_y, slim=(box_style == "sam-fast"))
                
                # Collect mask and point data
                all_masks.append(mask)
                point_coords.append((center_x, center_y))
                point_labels.append(detect_keyword)
                
            except Exception as e:
                print(f"Error processing individual SAM point: {str(e)}")
                print(f"Point data: {point}")
        
        if all_masks:
            # Create final visualization with all masks
            result_pil = create_mask_overlay(frame_pil, all_masks, point_coords, point_labels)
            frame = cv2.cvtColor(np.array(result_pil), cv2.COLOR_RGB2BGR)

    # Process other visualization styles
    for detection in detected_objects:
        try:
            # Handle both tracked and untracked detections
            if len(detection) == 3:  # Tracked detection with ID
                box, keyword, track_id = detection
            else:  # Regular detection without tracking
                box, keyword = detection
                track_id = None

            x1 = int(box[0] * width)
            y1 = int(box[1] * height)
            x2 = int(box[2] * width)
            y2 = int(box[3] * height)

            x1 = max(0, min(x1, width - 1))
            y1 = max(0, min(y1, height - 1))
            x2 = max(0, min(x2, width - 1))
            y2 = max(0, min(y2, height - 1))

            if x2 > x1 and y2 > y1:
                if box_style == "censor":
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 0), -1)
                elif box_style == "bounding-box":
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 3)

                    label = f"{detect_keyword}" if track_id is not None else detect_keyword
                    label_size = cv2.getTextSize(label, FONT, 0.7, 2)[0]
                    cv2.rectangle(
                        frame, (x1, y1 - 25), (x1 + label_size[0], y1), (0, 0, 255), -1
                    )
                    cv2.putText(
                        frame,
                        label,
                        (x1, y1 - 6),
                        FONT,
                        0.7,
                        (255, 255, 255),
                        2,
                        cv2.LINE_AA,
                    )
                elif box_style == "fuzzy-blur":
                    # Extract ROI
                    roi = frame[y1:y2, x1:x2]
                    # Apply Gaussian blur with much larger kernel for intense blur
                    blurred_roi = cv2.GaussianBlur(roi, (125, 125), 0)
                    # Replace original ROI with blurred version
                    frame[y1:y2, x1:x2] = blurred_roi
                elif box_style == "pixelated-blur":
                    # Extract ROI
                    roi = frame[y1:y2, x1:x2]
                    # Pixelate by resizing down and up
                    h, w = roi.shape[:2]
                    temp = cv2.resize(roi, (10, 10), interpolation=cv2.INTER_LINEAR)
                    pixelated = cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST)
                    # Mix up the pixelated frame slightly by adding random noise
                    noise = np.random.randint(0, 50, (h, w, 3), dtype=np.uint8)
                    pixelated = cv2.add(pixelated, noise)
                    # Apply stronger Gaussian blur to smooth edges
                    blurred_pixelated = cv2.GaussianBlur(pixelated, (15, 15), 0)
                    # Replace original ROI
                    frame[y1:y2, x1:x2] = blurred_pixelated
                elif box_style == "obfuscated-pixel":
                    # Calculate expansion amount based on 10% of object dimensions
                    box_width = x2 - x1
                    box_height = y2 - y1
                    expand_x = int(box_width * 0.10)
                    expand_y = int(box_height * 0.10)
                    
                    # Expand the bounding box by 10% in all directions
                    x1_expanded = max(0, x1 - expand_x)
                    y1_expanded = max(0, y1 - expand_y)
                    x2_expanded = min(width - 1, x2 + expand_x)
                    y2_expanded = min(height - 1, y2 + expand_y)
                    
                    # Extract ROI with much larger padding for true background sampling
                    padding = 100  # Much larger padding to get true background
                    y1_pad = max(0, y1_expanded - padding)
                    y2_pad = min(height, y2_expanded + padding)
                    x1_pad = max(0, x1_expanded - padding)
                    x2_pad = min(width, x2_expanded + padding)
                    
                    # Get the padded region including background
                    padded_roi = frame[y1_pad:y2_pad, x1_pad:x2_pad]
                    
                    # Create mask that excludes a larger region around the detection
                    h, w = y2_expanded - y1_expanded, x2_expanded - x1_expanded
                    bg_mask = np.ones(padded_roi.shape[:2], dtype=bool)
                    
                    # Exclude a larger region around the detection from background sampling
                    exclusion_padding = 50  # Area to exclude around detection
                    exclude_y1 = padding - exclusion_padding
                    exclude_y2 = padding + h + exclusion_padding
                    exclude_x1 = padding - exclusion_padding
                    exclude_x2 = padding + w + exclusion_padding
                    
                    # Make sure exclusion coordinates are valid
                    exclude_y1 = max(0, exclude_y1)
                    exclude_y2 = min(padded_roi.shape[0], exclude_y2)
                    exclude_x1 = max(0, exclude_x1)
                    exclude_x2 = min(padded_roi.shape[1], exclude_x2)
                    
                    # Mark the exclusion zone in the mask
                    bg_mask[exclude_y1:exclude_y2, exclude_x1:exclude_x2] = False
                    
                    # If we have enough background pixels, calculate average color
                    if np.any(bg_mask):
                        bg_color = np.mean(padded_roi[bg_mask], axis=0).astype(np.uint8)
                    else:
                        # Fallback to edges if we couldn't get enough background
                        edge_samples = np.concatenate([
                            padded_roi[0],  # Top edge
                            padded_roi[-1],  # Bottom edge
                            padded_roi[:, 0],  # Left edge
                            padded_roi[:, -1]  # Right edge
                        ])
                        bg_color = np.mean(edge_samples, axis=0).astype(np.uint8)
                    
                    # Create base pixelated version (of the expanded region)
                    temp = cv2.resize(frame[y1_expanded:y2_expanded, x1_expanded:x2_expanded], 
                                   (6, 6), interpolation=cv2.INTER_LINEAR)
                    pixelated = cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST)
                    
                    # Blend heavily towards background color
                    blend_factor = 0.9  # Much stronger blend with background
                    blended = cv2.addWeighted(
                        pixelated, 1 - blend_factor,
                        np.full((h, w, 3), bg_color, dtype=np.uint8), blend_factor,
                        0
                    )
                    
                    # Replace original ROI with blended version (using expanded coordinates)
                    frame[y1_expanded:y2_expanded, x1_expanded:x2_expanded] = blended
                elif box_style == "intense-pixelated-blur":
                    # Expand the bounding box by pixels in all directions
                    x1_expanded = max(0, x1 - 15)
                    y1_expanded = max(0, y1 - 15)
                    x2_expanded = min(width - 1, x2 + 25)
                    y2_expanded = min(height - 1, y2 + 25)

                    # Extract ROI
                    roi = frame[y1_expanded:y2_expanded, x1_expanded:x2_expanded]
                    # Pixelate by resizing down and up
                    h, w = roi.shape[:2]
                    temp = cv2.resize(roi, (10, 10), interpolation=cv2.INTER_LINEAR)
                    pixelated = cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST)
                    # Mix up the pixelated frame slightly by adding random noise
                    noise = np.random.randint(0, 50, (h, w, 3), dtype=np.uint8)
                    pixelated = cv2.add(pixelated, noise)
                    # Apply stronger Gaussian blur to smooth edges
                    blurred_pixelated = cv2.GaussianBlur(pixelated, (15, 15), 0)
                    # Replace original ROI
                    frame[y1_expanded:y2_expanded, x1_expanded:x2_expanded] = blurred_pixelated
                elif box_style == "hitmarker":
                    if points:
                        for point in points:
                            try:
                                print(f"Processing point: {point}")
                                center_x = int(float(point["x"]) * width)
                                center_y = int(float(point["y"]) * height)
                                print(f"Converted coordinates: ({center_x}, {center_y})")

                                draw_hitmarker(frame, center_x, center_y)

                                label = f"{detect_keyword}" if track_id is not None else detect_keyword
                                label_size = cv2.getTextSize(label, FONT, 0.5, 1)[0]
                                cv2.putText(
                                    frame,
                                    label,
                                    (center_x - label_size[0] // 2, center_y - HITMARKER_SIZE - 5),
                                    FONT,
                                    0.5,
                                    HITMARKER_COLOR,
                                    1,
                                    cv2.LINE_AA,
                                )
                            except Exception as e:
                                print(f"Error processing individual point: {str(e)}")
                                print(f"Point data: {point}")
                elif box_style == "magnify":
                    # Calculate the center point of the detection
                    center_x = (x1 + x2) // 2
                    center_y = (y1 + y2) // 2
                    
                    # Calculate original dimensions
                    orig_width = x2 - x1
                    orig_height = y2 - y1
                    
                    # Calculate new dimensions using magnify_factor parameter
                    magnify_factor = getattr(model, "magnify_factor", 2.0)  # Default to 2x if not specified
                    new_width = int(orig_width * magnify_factor)
                    new_height = int(orig_height * magnify_factor)
                    
                    # Calculate new coordinates ensuring they stay within frame bounds
                    new_x1 = max(0, center_x - new_width // 2)
                    new_y1 = max(0, center_y - new_height // 2)
                    new_x2 = min(width - 1, new_x1 + new_width)
                    new_y2 = min(height - 1, new_y1 + new_height)
                    
                    # Extract the original ROI
                    roi = frame[y1:y2, x1:x2]
                    
                    # Resize the ROI using the magnify_factor
                    enlarged_roi = cv2.resize(roi, (new_x2 - new_x1, new_y2 - new_y1))
                    
                    # Create a mask for smooth blending
                    mask = np.zeros((new_y2 - new_y1, new_x2 - new_x1), dtype=np.float32)
                    cv2.rectangle(mask, (0, 0), (new_x2 - new_x1, new_y2 - new_y1), 1, -1)
                    mask = cv2.GaussianBlur(mask, (21, 21), 11)
                    
                    # Blend the enlarged ROI with the original frame
                    for c in range(3):  # For each color channel
                        frame[new_y1:new_y2, new_x1:new_x2, c] = (
                            frame[new_y1:new_y2, new_x1:new_x2, c] * (1 - mask) +
                            enlarged_roi[:, :, c] * mask
                        )

        except Exception as e:
            print(f"Error drawing {box_style} style box: {str(e)}")
            print(f"Box data: {box}")
            print(f"Keyword: {keyword}")

    return frame


def filter_temporal_outliers(detections_dict):
    """Filter out extremely large detections that take up most of the frame.
    Only keeps detections that are reasonable in size.

    Args:
        detections_dict: Dictionary of {frame_number: [(box, keyword, track_id), ...]}
    """
    filtered_detections = {}

    for t, detections in detections_dict.items():
        # Only keep detections that aren't too large
        valid_detections = []
        for detection in detections:
            # Handle both tracked and untracked detections
            if len(detection) == 3:  # Tracked detection with ID
                box, keyword, track_id = detection
            else:  # Regular detection without tracking
                box, keyword = detection
                track_id = None

            # Calculate box size as percentage of frame
            width = box[2] - box[0]
            height = box[3] - box[1]
            area = width * height

            # If box is less than 90% of frame, keep it
            if area < 0.9:
                if track_id is not None:
                    valid_detections.append((box, keyword, track_id))
                else:
                    valid_detections.append((box, keyword))

        if valid_detections:
            filtered_detections[t] = valid_detections

    return filtered_detections


def describe_frames(video_path, model, tokenizer, detect_keyword, test_mode=False, test_duration=DEFAULT_TEST_MODE_DURATION, grid_rows=1, grid_cols=1):
    """Extract and detect objects in frames."""
    props = get_video_properties(video_path)
    fps = props["fps"]

    # Initialize DeepSORT tracker
    tracker = DeepSORTTracker()

    # If in test mode, only process first N seconds
    if test_mode:
        frame_count = min(int(fps * test_duration), props["frame_count"])
    else:
        frame_count = props["frame_count"]

    ad_detections = {}  # Store detection results by frame number

    print("Extracting frames and detecting objects...")
    video = cv2.VideoCapture(video_path)

    # Detect scenes first
    scenes = detect(video_path, scene_detector)
    scene_changes = set(end.get_frames() for _, end in scenes)
    print(f"Detected {len(scenes)} scenes")

    frame_count_processed = 0
    with tqdm(total=frame_count) as pbar:
        while frame_count_processed < frame_count:
            ret, frame = video.read()
            if not ret:
                break

            # Check if current frame is a scene change
            if frame_count_processed in scene_changes:
                # Detect objects in the frame
                detected_objects = detect_objects_in_frame(
                    model, tokenizer, frame, detect_keyword, grid_rows=grid_rows, grid_cols=grid_cols
                )

            # Update tracker with current detections
            tracked_objects = tracker.update(frame, detected_objects)

            # Store results for every frame, even if empty
            ad_detections[frame_count_processed] = tracked_objects

            frame_count_processed += 1
            pbar.update(1)

    video.release()

    if frame_count_processed == 0:
        print("No frames could be read from video")
        return {}

    return ad_detections


def create_detection_video(
    video_path,
    ad_detections,
    detect_keyword,
    model,
    output_path=None,
    ffmpeg_preset="medium",
    test_mode=False,
    test_duration=DEFAULT_TEST_MODE_DURATION,
    box_style="censor",
):
    """Create video with detection boxes while preserving audio."""
    if output_path is None:
        # Create outputs directory if it doesn't exist
        outputs_dir = os.path.join(
            os.path.dirname(os.path.abspath(__file__)), "outputs"
        )
        os.makedirs(outputs_dir, exist_ok=True)

        # Clean the detect_keyword for filename
        safe_keyword = "".join(
            x for x in detect_keyword if x.isalnum() or x in (" ", "_", "-")
        )
        safe_keyword = safe_keyword.replace(" ", "_")

        # Create output filename
        base_name = os.path.splitext(os.path.basename(video_path))[0]
        output_path = os.path.join(
            outputs_dir, f"{box_style}_{safe_keyword}_{base_name}.mp4"
        )

    print(f"Will save output to: {output_path}")

    props = get_video_properties(video_path)
    fps, width, height = props["fps"], props["width"], props["height"]

    # If in test mode, only process first few seconds
    if test_mode:
        frame_count = min(int(fps * test_duration), props["frame_count"])
        print(f"Test mode enabled: Processing first {test_duration} seconds ({frame_count} frames)")
    else:
        frame_count = props["frame_count"]
        print("Full video mode: Processing entire video")

    video = cv2.VideoCapture(video_path)

    # Create temp output path by adding _temp before the extension
    base, ext = os.path.splitext(output_path)
    temp_output = f"{base}_temp{ext}"
    temp_audio = f"{base}_audio.aac"  # Temporary audio file

    out = cv2.VideoWriter(
        temp_output, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)
    )

    print("Creating detection video...")
    frame_count_processed = 0

    with tqdm(total=frame_count) as pbar:
        while frame_count_processed < frame_count:
            ret, frame = video.read()
            if not ret:
                break

            # Get detections for this exact frame
            if frame_count_processed in ad_detections:
                current_detections = ad_detections[frame_count_processed]
                if current_detections:
                    frame = draw_ad_boxes(
                        frame, current_detections, detect_keyword, model, box_style=box_style
                    )

            out.write(frame)
            frame_count_processed += 1
            pbar.update(1)

    video.release()
    out.release()

    # Extract audio from original video
    try:
        if test_mode:
            # In test mode, extract only the required duration of audio
            subprocess.run(
                [
                    "ffmpeg",
                    "-y",
                    "-i",
                    video_path,
                    "-t",
                    str(test_duration),
                    "-vn",  # No video
                    "-acodec",
                    "copy",
                    temp_audio,
                ],
                check=True,
            )
        else:
            subprocess.run(
                [
                    "ffmpeg",
                    "-y",
                    "-i",
                    video_path,
                    "-vn",  # No video
                    "-acodec",
                    "copy",
                    temp_audio,
                ],
                check=True,
            )
    except subprocess.CalledProcessError as e:
        print(f"Error extracting audio: {str(e)}")
        if os.path.exists(temp_output):
            os.remove(temp_output)
        return None

    # Merge processed video with original audio
    try:
        # Base FFmpeg command
        ffmpeg_cmd = [
            "ffmpeg",
            "-y",
            "-i",
            temp_output,
            "-i",
            temp_audio,
            "-c:v",
            "libx264",
            "-preset",
            ffmpeg_preset,
            "-crf",
            "23",
            "-c:a",
            "aac",
            "-b:a",
            "192k",
            "-movflags",
            "+faststart",  # Better web playback
        ]

        if test_mode:
            # In test mode, ensure output duration matches test_duration
            ffmpeg_cmd.extend([
                "-t",
                str(test_duration),
                "-shortest"  # Ensure output duration matches shortest input
            ])

        ffmpeg_cmd.extend([
            "-loglevel",
            "error",
            output_path
        ])

        subprocess.run(ffmpeg_cmd, check=True)

        # Clean up temporary files
        os.remove(temp_output)
        os.remove(temp_audio)

        if not os.path.exists(output_path):
            print(
                f"Warning: FFmpeg completed but output file not found at {output_path}"
            )
            return None

        return output_path

    except subprocess.CalledProcessError as e:
        print(f"Error merging audio with video: {str(e)}")
        if os.path.exists(temp_output):
            os.remove(temp_output)
        if os.path.exists(temp_audio):
            os.remove(temp_audio)
        return None


def process_video(
    video_path,
    target_object,
    test_mode=False,
    test_duration=DEFAULT_TEST_MODE_DURATION,
    ffmpeg_preset="medium",
    grid_rows=1,
    grid_cols=1,
    box_style="censor",
    magnify_factor=2.0,
):
    """Process a video to detect and visualize specified objects."""
    try:
        print(f"\nProcessing: {video_path}")
        print(f"Looking for: {target_object}")

        # Load model
        print("Loading Moondream model...")
        model, tokenizer = load_moondream()
        
        # Add magnify_factor to model dict for use in draw_ad_boxes
        model.magnify_factor = magnify_factor

        # Get video properties
        props = get_video_properties(video_path)
        
        # Initialize scene detector with ContentDetector
        scene_detector = ContentDetector(threshold=30.0)  # Adjust threshold as needed
        
        # Initialize DeepSORT tracker
        tracker = DeepSORTTracker()

        # If in test mode, only process first N seconds
        if test_mode:
            frame_count = min(int(props["fps"] * test_duration), props["frame_count"])
        else:
            frame_count = props["frame_count"]

        ad_detections = {}  # Store detection results by frame number

        print("Extracting frames and detecting objects...")
        video = cv2.VideoCapture(video_path)

        # Detect scenes first
        scenes = detect(video_path, scene_detector)
        scene_changes = set(end.get_frames() for _, end in scenes)
        print(f"Detected {len(scenes)} scenes")

        frame_count_processed = 0
        with tqdm(total=frame_count) as pbar:
            while frame_count_processed < frame_count:
                ret, frame = video.read()
                if not ret:
                    break

                # Check if current frame is a scene change
                if frame_count_processed in scene_changes:
                    print(f"Scene change detected at frame {frame_count_processed}. Resetting tracker.")
                    tracker.reset()

                # Detect objects in the frame
                detected_objects = detect_objects_in_frame(
                    model, tokenizer, frame, target_object, grid_rows=grid_rows, grid_cols=grid_cols
                )

                # Update tracker with current detections
                tracked_objects = tracker.update(frame, detected_objects)

                # Store results for every frame, even if empty
                ad_detections[frame_count_processed] = tracked_objects

                frame_count_processed += 1
                pbar.update(1)

        video.release()

        if frame_count_processed == 0:
            print("No frames could be read from video")
            return {}

        # Apply filtering
        filtered_ad_detections = filter_temporal_outliers(ad_detections)
        
        # Build detection data structure
        detection_data = {
            "video_metadata": {
                "file_name": os.path.basename(video_path),
                "fps": props["fps"],
                "width": props["width"],
                "height": props["height"],
                "total_frames": props["frame_count"],
                "duration_sec": props["frame_count"] / props["fps"],
                "detect_keyword": target_object,
                "test_mode": test_mode,
                "grid_size": f"{grid_rows}x{grid_cols}",
                "box_style": box_style,
                "timestamp": datetime.now().isoformat()
            },
            "frame_detections": [
                {
                    "frame": frame_num,
                    "timestamp": frame_num / props["fps"],
                    "objects": [
                        {
                            "keyword": kw,
                            "bbox": list(box),  # Convert numpy array to list if needed
                            "track_id": track_id if len(detection) == 3 else None
                        }
                        for detection in filtered_ad_detections.get(frame_num, [])
                        for box, kw, *track_id in [detection]  # Unpack detection tuple, track_id will be empty list if not present
                    ]
                }
                for frame_num in range(props["frame_count"] if not test_mode else min(int(props["fps"] * test_duration), props["frame_count"]))
            ]
        }
        
        # Save filtered data
        outputs_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputs")
        os.makedirs(outputs_dir, exist_ok=True)
        base_name = os.path.splitext(os.path.basename(video_path))[0]
        json_path = os.path.join(outputs_dir, f"{box_style}_{target_object}_{base_name}_detections.json")
        
        from persistence import save_detection_data
        if not save_detection_data(detection_data, json_path):
            print("Warning: Failed to save detection data")

        # Create video with filtered data
        output_path = create_detection_video(
            video_path,
            filtered_ad_detections,
            target_object,
            model,
            ffmpeg_preset=ffmpeg_preset,
            test_mode=test_mode,
            test_duration=test_duration,
            box_style=box_style,
        )

        if output_path is None:
            print("\nError: Failed to create output video")
            return None

        print(f"\nOutput saved to: {output_path}")
        print(f"Detection data saved to: {json_path}")
        return output_path
        
    except Exception as e:
        print(f"Error processing video: {str(e)}")
        import traceback
        traceback.print_exc()
        return None


def main():
    """Process all videos in the inputs directory."""
    parser = argparse.ArgumentParser(
        description="Detect objects in videos using Moondream2"
    )
    parser.add_argument(
        "--test", action="store_true", help="Process only first 3 seconds of each video"
    )
    parser.add_argument(
        "--test-duration",
        type=int,
        default=DEFAULT_TEST_MODE_DURATION,
        help=f"Number of seconds to process in test mode (default: {DEFAULT_TEST_MODE_DURATION})"
    )
    parser.add_argument(
        "--preset",
        choices=FFMPEG_PRESETS,
        default="medium",
        help="FFmpeg encoding preset (default: medium). Faster presets = lower quality",
    )
    parser.add_argument(
        "--detect",
        type=str,
        default="face",
        help='Object to detect in the video (default: face, use --detect "thing to detect" to override)',
    )
    parser.add_argument(
        "--rows",
        type=int,
        default=1,
        help="Number of rows to split each frame into (default: 1)",
    )
    parser.add_argument(
        "--cols",
        type=int,
        default=1,
        help="Number of columns to split each frame into (default: 1)",
    )
    parser.add_argument(
        "--box-style",
        choices=["censor", "bounding-box", "hitmarker", "sam", "sam-fast", "fuzzy-blur", 
                "pixelated-blur", "intense-pixelated-blur", "obfuscated-pixel", "magnify"],
        default="censor",
        help="Style of detection visualization (default: censor)",
    )
    args = parser.parse_args()

    input_dir = "inputs"
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs("outputs", exist_ok=True)

    video_files = [
        f
        for f in os.listdir(input_dir)
        if f.lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm"))
    ]

    if not video_files:
        print("No video files found in 'inputs' directory")
        return

    print(f"Found {len(video_files)} videos to process")
    print(f"Will detect: {args.detect}")
    if args.test:
        print("Running in test mode - processing only first 3 seconds of each video")
    print(f"Using FFmpeg preset: {args.preset}")
    print(f"Grid size: {args.rows}x{args.cols}")
    print(f"Box style: {args.box_style}")

    success_count = 0
    for video_file in video_files:
        video_path = os.path.join(input_dir, video_file)
        output_path = process_video(
            video_path,
            args.detect,
            test_mode=args.test,
            test_duration=args.test_duration,
            ffmpeg_preset=args.preset,
            grid_rows=args.rows,
            grid_cols=args.cols,
            box_style=args.box_style,
            magnify_factor=args.magnify_factor,
        )
        if output_path:
            success_count += 1

    print(
        f"\nProcessing complete. Successfully processed {success_count} out of {len(video_files)} videos."
    )


if __name__ == "__main__":
    main()