app.py

import gradio as gr
import cv2
import numpy as np
import torch
import os
import tempfile
import subprocess

import torch
model = torch.hub.load('ultralytics/yolov5', 'custom', path='watermark-detection-73.pt', force_reload=True)

model.conf = 0.25
model.iou = 0.45
model.max_det = 1


# ------------------------------
# Helper Functions
# ------------------------------
def extract_first_frame(video_path):
    cap = cv2.VideoCapture(video_path)
    ret, frame = cap.read()
    cap.release()
    if ret:
        return frame
    return None

def detect_watermark_coordinates(frame):
    results = model(frame)
    detections = results.xyxy[0].cpu().numpy()
    if len(detections) == 0:
        return None
    x1, y1, x2, y2, _, _ = detections[0]
    return int(x1), int(y1), int(x2 - x1), int(y2 - y1)

def generate_mask_from_coords(frame_shape, x, y, w, h):
    mask = np.zeros(frame_shape[:2], dtype=np.uint8)
    mask[int(y):int(y+h), int(x):int(x+w)] = 255
    return mask

def apply_inpaint_to_video(video_path, x, y, w, h):
    temp_dir = tempfile.mkdtemp()
    cap = cv2.VideoCapture(video_path)
    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))
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    output_video_path = os.path.join(temp_dir, "output.mp4")
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
    for _ in range(frame_count):
        ret, frame = cap.read()
        if not ret:
            break
        mask = generate_mask_from_coords(frame.shape, x, y, w, h)
        inpainted = cv2.inpaint(frame, mask, 3, cv2.INPAINT_TELEA)
        writer.write(inpainted)
    cap.release()
    writer.release()

    # Combine processed video with original audio
    temp_with_audio = os.path.splitext(video_path)[0] + "_no_watermark.mp4"
    cmd_audio = f'ffmpeg -y -i "{output_video_path}" -i "{video_path}" -c:v copy -c:a aac -map 0:v:0 -map 1:a:0 -shortest "{temp_with_audio}"'
    subprocess.call(cmd_audio, shell=True)

    # Re-encode to ensure browser compatibility
    final_output_path = os.path.splitext(video_path)[0] + "_no_watermark_fixed.mp4"
    cmd_fix = f'ffmpeg -y -i "{temp_with_audio}" -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" -c:v libx264 -preset fast -crf 23 -c:a aac -b:a 128k -movflags +faststart "{final_output_path}"'
    subprocess.call(cmd_fix, shell=True)
    return final_output_path

def apply_inpaint_to_image(image, x, y, w, h):
    mask = generate_mask_from_coords(image.shape, x, y, w, h)
    inpainted = cv2.inpaint(image, mask, 3, cv2.INPAINT_TELEA)
    return inpainted

def overlay_box_on_image(image, x, y, w, h):
    image_with_box = image.copy()
    cv2.rectangle(image_with_box, (x, y), (x + w, y + h), (0, 255, 0), 2)
    return image_with_box

def get_coords_for_image(image):
    coords = detect_watermark_coordinates(image)
    h, w, _ = image.shape
    if coords:
        x, y, w_box, h_box = coords
        return overlay_box_on_image(image, x, y, w_box, h_box), x, y, w_box, h_box, "✅ Auto watermark detected."
    else:
        w_box, h_box = int(w * 0.25), int(h * 0.1)
        x, y = (w - w_box) // 2, (h - h_box) // 2
        return overlay_box_on_image(image, x, y, w_box, h_box), x, y, w_box, h_box, "⚠️ No watermark detected. Default box placed."

def update_image_live(image, x, y, w, h):
    return overlay_box_on_image(image, int(x), int(y), int(w), int(h))

def process_uploaded_video(video, x, y, w, h):
    try:
        output_path = apply_inpaint_to_video(video, int(x), int(y), int(w), int(h))
        return output_path, "✅ Watermark removed from video."
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def process_uploaded_image(image, x, y, w, h):
    try:
        result = apply_inpaint_to_image(image, int(x), int(y), int(w), int(h))
        return result, "✅ Watermark removed from image."
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# ------------------------------
# Gradio UI (Merged with Theme)
# ------------------------------
with gr.Blocks(theme=gr.themes.Soft(), title="Watermark Remover") as demo:
    gr.Markdown("<p style='text-align: center;'>Remove watermarks from both videos and images using AI detection or manual box selection.</p>")

    with gr.Tab("📹 Video Watermark Remover"):
        with gr.Row():
            with gr.Column(scale=1):
                video_input = gr.Video(label="🎞️ Upload Video")
                auto_btn_v = gr.Button("🔍 Auto Detect Watermark", variant="primary")
                run_btn_v = gr.Button("🧹 Remove Watermark", variant="secondary")
                status_v = gr.Textbox(label="Status", interactive=False)
                output_file_v = gr.File(label="⬇️ Download Cleaned Video")
            with gr.Column(scale=1):
                video_frame = gr.Image(label="📍 Watermark Preview", interactive=False)
                frame_original = gr.State()
                with gr.Accordion("🔧 Manual Box Adjustment", open=False):
                    x_v = gr.Slider(minimum=0, maximum=2000, label="X Coordinate", step=1)
                    y_v = gr.Slider(minimum=0, maximum=2000, label="Y Coordinate", step=1)
                    w_v = gr.Slider(minimum=10, maximum=2000, label="Width", step=1)
                    h_v = gr.Slider(minimum=10, maximum=2000, label="Height", step=1)

        auto_btn_v.click(
            fn=lambda video: get_coords_for_image(extract_first_frame(video)),
            inputs=video_input,
            outputs=[video_frame, x_v, y_v, w_v, h_v, status_v],
        ).then(
            fn=lambda video: extract_first_frame(video),
            inputs=video_input,
            outputs=frame_original,
        )

        for slider in [x_v, y_v, w_v, h_v]:
            slider.change(
                fn=update_image_live,
                inputs=[frame_original, x_v, y_v, w_v, h_v],
                outputs=video_frame
            )

        run_btn_v.click(
            fn=process_uploaded_video,
            inputs=[video_input, x_v, y_v, w_v, h_v],
            outputs=[output_file_v, status_v]
        )

    with gr.Tab("🖼️ Image Watermark Remover"):
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(label="🖼️ Upload Image")
                auto_btn_i = gr.Button("🔍 Auto Detect Watermark", variant="primary")
                run_btn_i = gr.Button("🧹 Remove Watermark", variant="secondary")
                status_i = gr.Textbox(label="Status", interactive=False)
                output_image = gr.Image(label="🧼 Cleaned Image")
            with gr.Column(scale=1):
                image_display = gr.Image(label="📍 Watermark Preview", interactive=False)
                image_original = gr.State()
                with gr.Accordion("🔧 Manual Box Adjustment", open=False):
                    x_i = gr.Slider(minimum=0, maximum=2000, label="X Coordinate", step=1)
                    y_i = gr.Slider(minimum=0, maximum=2000, label="Y Coordinate", step=1)
                    w_i = gr.Slider(minimum=10, maximum=2000, label="Width", step=1)
                    h_i = gr.Slider(minimum=10, maximum=2000, label="Height", step=1)

        auto_btn_i.click(
            fn=get_coords_for_image,
            inputs=image_input,
            outputs=[image_display, x_i, y_i, w_i, h_i, status_i],
        ).then(
            fn=lambda img: img,
            inputs=image_input,
            outputs=image_original,
        )

        for slider in [x_i, y_i, w_i, h_i]:
            slider.change(
                fn=update_image_live,
                inputs=[image_original, x_i, y_i, w_i, h_i],
                outputs=image_display
            )

        run_btn_i.click(
            fn=process_uploaded_image,
            inputs=[image_input, x_i, y_i, w_i, h_i],
            outputs=[output_image, status_i]
        )

if __name__ == '__main__':
    demo.launch()