Create app.py

app.py

@@ -0,0 +1,197 @@
+import gradio as gr
+import cv2
+import numpy as np
+import torch
+import os
+import tempfile
+import subprocess
+
+from ultralytics import YOLO
+
+model = YOLO("watermark-detection-73.pt")
+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()