main_video.py

import cv2
import numpy as np
import os
import shutil
import subprocess
import urllib.request
import onnxruntime as ort
from moviepy.editor import VideoFileClip

# 🚀 INITIALIZE THE LAMA AI ENGINE GLOBALLY
MODEL_PATH = os.path.join(os.path.dirname(__file__), "models", "big-lama.onnx")

if not os.path.exists(MODEL_PATH):
    print("📥 AI Model file not found. Auto-downloading from production mirror...")
    os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
    url = "https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx"
    try:
        req = urllib.request.Request(
            url, 
            headers={'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)'}
        )
        with urllib.request.urlopen(req) as response, open(MODEL_PATH, 'wb') as out_file:
            shutil.copyfileobj(response, out_file)
        print("✅ AI Model downloaded successfully!")
    except Exception as download_err:
        print(f"🚨 Mirror download failed ({download_err}).")

try:
    if os.path.exists(MODEL_PATH):
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = 1
        opts.inter_op_num_threads = 1
        opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        ort_session = ort.InferenceSession(MODEL_PATH, sess_options=opts, providers=providers)
        print("🚀 LaMa AI Inpainting Engine Initialized Safely!")
    else:
        ort_session = None
except Exception as e:
    print(f"⚠️ ONNX Session initialization failed ({e}). Falling back to classical.")
    ort_session = None


def remove_watermark_pro(input_path, output_path, x, y, w, h, mode="fast"):
    """
    Video Watermark Removal with adaptive dominant-color patching for seamless graphics removal.
    """
    global ort_session
    
    base_name = os.path.basename(input_path)
    normalized_input = f"normalized_{base_name}.mp4"
    
    try:
        conversion_cmd = [
            "ffmpeg", "-y", "-i", input_path,
            "-vcodec", "libx264", "-acodec", "aac",
            "-pix_fmt", "yuv420p", "-preset", "ultrafast",
            normalized_input
        ]
        subprocess.run(conversion_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        processing_source = normalized_input
    except Exception:
        processing_source = input_path

    cap = cv2.VideoCapture(processing_source)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    width = width if width % 2 == 0 else width - 1
    height = height if height % 2 == 0 else height - 1

    temp_output = f"temp_pro_{base_name}"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))

    # Calculate bounded crop area coordinates for AI mode
    padding = 16
    x1 = max(0, int(x) - padding)
    y1 = max(0, int(y) - padding)
    x2 = min(width, int(x + w) + padding)
    y2 = min(height, int(y + h) + padding)

    crop_w = x2 - x1
    crop_h = y2 - y1
    if crop_w % 8 != 0: x2 = min(width, x2 + (8 - (crop_w % 8)))
    if crop_h % 8 != 0: y2 = min(height, y2 + (8 - (crop_h % 8)))

    img_key, mask_key = None, None
    if ort_session is not None and mode == "ai":
        try:
            inputs = ort_session.get_inputs()
            for node in inputs:
                if "mask" in node.name.lower():
                    mask_key = node.name
                else:
                    img_key = node.name
        except Exception:
            ort_session = None

    frames_processed = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        if frame.shape[1] != width or frame.shape[0] != height:
            frame = cv2.resize(frame, (width, height))
            
        mask = np.zeros(frame.shape[:2], dtype=np.uint8)
        cv2.rectangle(mask, (int(x), int(y)), (int(x + w), int(y + h)), 255, -1)
        
        if mode == "ai" and ort_session is not None and (y2 > y1 and x2 > x1):
            try:
                crop_frame = frame[y1:y2, x1:x2]
                crop_mask = mask[y1:y2, x1:x2]
                orig_h, orig_w = crop_frame.shape[:2]

                crop_frame_resized = cv2.resize(crop_frame, (512, 512))
                crop_mask_resized = cv2.resize(crop_mask, (512, 512))
                crop_rgb = cv2.cvtColor(crop_frame_resized, cv2.COLOR_BGR2RGB)
                
                img_tensor = crop_rgb.astype(np.float32) / 255.0
                img_tensor = np.transpose(img_tensor, (2, 0, 1))[np.newaxis, ...]
                
                mask_tensor = crop_mask_resized.astype(np.float32)
                mask_tensor = np.where(mask_tensor > 0, 1.0, 0.0).astype(np.float32)
                mask_tensor = mask_tensor[np.newaxis, np.newaxis, ...]

                ai_outputs = ort_session.run(None, {img_key: img_tensor, mask_key: mask_tensor})

                ai_output = ai_outputs[0][0]
                ai_output = np.transpose(ai_output, (1, 2, 0))
                ai_output = np.clip(ai_output * 255.0, 0, 255).astype(np.uint8)
                clean_crop_bgr = cv2.cvtColor(ai_output, cv2.COLOR_RGB2BGR)

                frame[y1:y2, x1:x2] = cv2.resize(clean_crop_bgr, (orig_w, orig_h))
            except Exception:
                frame = cv2.inpaint(frame, mask, 1, cv2.INPAINT_TELEA)
        else:
            # ⚡ HIGH-FIDELITY FAST MODE: Content-Aware Dominant Color Patching
            ymin, ymax = max(0, int(y)), min(height, int(y + h))
            xmin, xmax = max(0, int(x)), min(width, int(x + w))
            
            if (ymax > ymin) and (xmax > xmin):
                # Sample 4 pixels wide around the boundary of the selection box
                sample_top = frame[max(0, ymin-4):ymin, xmin:xmax]
                sample_bottom = frame[ymax:min(height, ymax+4), xmin:xmax]
                sample_left = frame[ymin:ymax, max(0, xmin-4):xmin]
                sample_right = frame[ymin:ymax, xmax:min(width, xmax+4)]
                
                samples = []
                for s in [sample_top, sample_bottom, sample_left, sample_right]:
                    if s.size > 0:
                        samples.append(s.reshape(-1, 3))
                
                if samples:
                    # Use median value to completely disregard text letter colors and extract pure background tint
                    all_pixels = np.vstack(samples)
                    bg_color = np.median(all_pixels, axis=0).astype(int)
                    
                    # Paint a clean, solid patch over the watermark zone matching the background perfectly
                    cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (int(bg_color[0]), int(bg_color[1]), int(bg_color[2])), -1)
                    
                    # Smooth the local border bounds slightly so it blends organically
                    if xmin > 2 and xmax < width - 2 and ymin > 2 and ymax < height - 2:
                        roi = frame[ymin-2:ymax+2, xmin-2:xmax+2]
                        frame[ymin-2:ymax+2, xmin-2:xmax+2] = cv2.GaussianBlur(roi, (3, 3), 0)
                else:
                    frame = cv2.inpaint(frame, mask, 1, cv2.INPAINT_TELEA)
            else:
                frame = cv2.inpaint(frame, mask, 1, cv2.INPAINT_TELEA)

        out.write(frame)
        frames_processed += 1

    cap.release()
    out.release()

    if frames_processed == 0:
        if os.path.exists(temp_output): os.remove(temp_output)
        if os.path.exists(normalized_input): os.remove(normalized_input)
        return False 

    try:
        original_clip = VideoFileClip(processing_source)
        processed_clip = VideoFileClip(temp_output)
        
        if original_clip.audio is not None:
            final_video = processed_clip.set_audio(original_clip.audio)
        else:
            final_video = processed_clip
            
        final_video.write_videofile(
            output_path, codec="libx264", audio_codec="aac", 
            preset="ultrafast", threads=1,
            ffmpeg_params=["-movflags", "+faststart"], logger=None
        )
        original_clip.close()
        processed_clip.close()
    except Exception:
        try:
            conversion_command = [
                "ffmpeg", "-y", "-i", temp_output, 
                "-vcodec", "libx264", "-pix_fmt", "yuv420p",
                "-movflags", "+faststart", "-an", output_path
            ]
            subprocess.run(conversion_command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except Exception:
            shutil.copy(temp_output, output_path)
    finally:
        if os.path.exists(temp_output): os.remove(temp_output)
        if os.path.exists(normalized_input): os.remove(normalized_input)

    return True