main.py

#!/usr/bin/env python3
import cv2, os, subprocess, argparse
from PIL import Image
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from tqdm import tqdm
import numpy as np
from datetime import datetime

# Constants
TEST_MODE_DURATION = 3  # Process only first 3 seconds in test mode
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

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_box(box):
    """Check if box coordinates are reasonable."""
    x1, y1, x2, y2 = 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_tiles(frame, rows, cols):
    """Split a frame into a grid of tiles."""
    height, width = frame.shape[:2]
    tile_height = height // rows
    tile_width = width // cols
    tiles = []
    tile_positions = []
    
    for i in range(rows):
        for j in range(cols):
            y1 = i * tile_height
            y2 = (i + 1) * tile_height if i < rows - 1 else height
            x1 = j * tile_width
            x2 = (j + 1) * tile_width if j < 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_ads_in_frame(model, tokenizer, image, detect_keyword, rows=1, cols=1):
    """Detect objects in a frame using grid-based detection."""
    if rows == 1 and cols == 1:
        return detect_ads_in_frame_single(model, tokenizer, image, detect_keyword)
    
    # 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_tiles(image, rows, 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, detect_keyword)
        
        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_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, detect_keyword))
            
            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_ads_in_frame_single(model, tokenizer, image, detect_keyword):
    """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, detect_keyword)
    
    # 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_box(box):
                    detected_objects.append((box, detect_keyword))
    
    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, box_style='censor'):
    """Draw detection visualizations over detected objects.
    
    Args:
        frame: The video frame to draw on
        detected_objects: List of (box, keyword) tuples
        detect_keyword: The detection keyword
        box_style: Visualization style ('censor', 'bounding-box', or 'hitmarker')
    """
    height, width = frame.shape[:2]
    
    for (box, keyword) in detected_objects:
        try:
            # Convert normalized coordinates to pixel coordinates
            x1 = int(box[0] * width)
            y1 = int(box[1] * height)
            x2 = int(box[2] * width)
            y2 = int(box[3] * height)
            
            # Ensure coordinates are within frame boundaries
            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))
            
            # Only draw if box has reasonable size
            if x2 > x1 and y2 > y1:
                if box_style == 'censor':
                    # Draw solid black rectangle
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 0), -1)
                elif box_style == 'bounding-box':
                    # Draw red rectangle with thicker line
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 3)
                    
                    # Add label with background
                    label = detect_keyword  # Use exact capitalization
                    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 == 'hitmarker':
                    # Calculate center of the box
                    center_x = (x1 + x2) // 2
                    center_y = (y1 + y2) // 2
                    
                    # Draw hitmarker at the center
                    draw_hitmarker(frame, center_x, center_y)
                    
                    # Optional: Add small label above hitmarker
                    label = detect_keyword  # Use exact capitalization
                    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 drawing {box_style} style box: {str(e)}")
    
    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), ...]}
    """
    filtered_detections = {}
    
    for t, detections in detections_dict.items():
        # Only keep detections that aren't too large
        valid_detections = []
        for box, keyword in detections:
            # 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:
                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, rows=1, cols=1):
    """Extract and detect objects in frames."""
    props = get_video_properties(video_path)
    fps = props['fps']
    
    # If in test mode, only process first 3 seconds
    if test_mode:
        frame_count = min(int(fps * TEST_MODE_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)
    
    # Process every frame
    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
            
            # Detect objects in the frame
            detected_objects = detect_ads_in_frame(model, tokenizer, frame, detect_keyword, rows=rows, cols=cols)
            
            # Store results for every frame, even if empty
            ad_detections[frame_count_processed] = detected_objects
            
            frame_count_processed += 1
            pbar.update(1)
    
    video.release()
    
    if frame_count_processed == 0:
        print("No frames could be read from video")
        return {}
    
    # Filter out only extremely large detections
    ad_detections = filter_temporal_outliers(ad_detections)
    return ad_detections

def create_detection_video(video_path, ad_detections, detect_keyword, output_path=None, ffmpeg_preset='medium', test_mode=False, box_style='censor'):
    """Create video with detection boxes."""
    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_MODE_DURATION), props['frame_count'])
    else:
        frame_count = props['frame_count']
    
    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}"
    
    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, box_style=box_style)
            
            out.write(frame)
            frame_count_processed += 1
            pbar.update(1)
    
    video.release()
    out.release()
    
    # Convert to web-compatible format more efficiently
    try:
        subprocess.run([
            'ffmpeg', '-y',
            '-i', temp_output,
            '-c:v', 'libx264',
            '-preset', ffmpeg_preset,
            '-crf', '23',
            '-movflags', '+faststart',  # Better web playback
            '-loglevel', 'error',
            output_path
        ], check=True)
        
        os.remove(temp_output)  # Remove the temporary file
        
        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 running FFmpeg: {str(e)}")
        if os.path.exists(temp_output):
            os.remove(temp_output)
        return None

def process_video(video_path, detect_keyword, test_mode=False, ffmpeg_preset='medium', rows=1, cols=1, box_style='censor'):
    """Process a single video file."""
    print(f"\nProcessing: {video_path}")
    print(f"Looking for: {detect_keyword}")
    
    # Load model
    print("Loading Moondream model...")
    model, tokenizer = load_moondream()
    
    # Process video - detect objects
    ad_detections = describe_frames(video_path, model, tokenizer, detect_keyword, test_mode, rows, cols)
    
    # Create video with detection boxes
    output_path = create_detection_video(video_path, ad_detections, detect_keyword, 
                                       ffmpeg_preset=ffmpeg_preset, test_mode=test_mode,
                                       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}")
    return output_path

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('--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'], 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, ffmpeg_preset=args.preset,
                     rows=args.rows, cols=args.cols, box_style=args.box_style)
        if output_path:
            success_count += 1
    
    print(f"\nProcessing complete. Successfully processed {success_count} out of {len(video_files)} videos.")

if __name__ == "__main__":
    main()