detect.py

import cv2
import numpy as np
import random
import tempfile
from moviepy.video.io.VideoFileClip import VideoFileClip

def detect_watermark_image(image):
    ycrcb_image = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
    y_channel, _, _ = cv2.split(ycrcb_image)
    dct_y = cv2.dct(np.float32(y_channel))
    
    # Detecting the watermark
    watermark = np.zeros_like(dct_y)
    rows, cols = dct_y.shape
    font = cv2.FONT_HERSHEY_SIMPLEX
    text = "WATERMARK"
    text_size = cv2.getTextSize(text, font, 0.5, 1)[0]
    text_x = np.random.randint(0, cols - text_size[0])
    text_y = np.random.randint(text_size[1], rows)
    watermark = cv2.putText(watermark, text, (text_x, text_y), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
    
    detected_image = cv2.idct(dct_y + watermark)
    detected_image = np.uint8(np.clip(detected_image, 0, 255))
    
    return detected_image

def detect_watermark_video(video_path):
    """Detect and highlight watermarks in a video file.
    
    Args:
        video_path (str): Path to the video file
        
    Returns:
        str: Path to the video with detected watermarks
    """
    def process_frame(frame):
        ycrcb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2YCrCb)
        y_channel, _, _ = cv2.split(ycrcb_image)
        dct_y = cv2.dct(np.float32(y_channel))
        
        # Detecting the watermark
        watermark = np.zeros_like(dct_y)
        rows, cols = dct_y.shape
        font = cv2.FONT_HERSHEY_SIMPLEX
        text = "WATERMARK"
        text_size = cv2.getTextSize(text, font, 0.5, 1)[0]
        text_x = np.random.randint(0, cols - text_size[0])
        text_y = np.random.randint(text_size[1], rows)
        watermark = cv2.putText(watermark, text, (text_x, text_y), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
        
        detected_frame = cv2.idct(dct_y + watermark)
        detected_frame = np.uint8(np.clip(detected_frame, 0, 255))
        
        # Reconstruct the image with detected watermark
        ycrcb_image[:, :, 0] = detected_frame
        output_frame = cv2.cvtColor(ycrcb_image, cv2.COLOR_YCrCb2BGR)
        
        return output_frame
    
    # Process the video
    video = moviepy.VideoFileClip(video_path)
    processed_video = video.fl_image(process_frame)
    
    # Save the result to a temporary file
    _, output_path = tempfile.mkstemp(suffix=".mp4")
    processed_video.write_videofile(output_path, codec='libx264')
    
    return output_path