detect.py

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

def detect_watermark_image(image, watermark_text="WATERMARK", threshold=0.7):
    """Detect if an image contains a DCT-domain watermark.
    
    Args:
        image: Input image
        watermark_text: The text to look for
        threshold: Detection sensitivity threshold
        
    Returns:
        tuple: (detected (bool), highlighted_image)
    """
    # Ensure image dimensions are divisible by 8 for DCT
    h, w = image.shape[:2]
    h_new = (h // 8) * 8
    w_new = (w // 8) * 8
    image = cv2.resize(image, (w_new, h_new))
    
    # Convert to YCrCb and extract Y channel
    ycrcb_image = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
    y_channel, cr, cb = cv2.split(ycrcb_image)
    
    # Apply DCT
    dct_y = cv2.dct(np.float32(y_channel))
    
    # Create known pattern for detection (systematic approach)
    h_blocks, w_blocks = dct_y.shape[0] // 8, dct_y.shape[1] // 8
    detection_score = 0
    max_location = None
    highlighted_image = image.copy()
    
    # Search for watermark pattern in blocks
    for y_block in range(h_blocks):
        for x_block in range(w_blocks):
            # Extract block
            block = dct_y[y_block*8:(y_block+1)*8, x_block*8:(x_block+1)*8]
            
            # Calculate variance to detect unusual patterns
            block_variance = np.var(block)
            # Check for anomalies in mid-frequency coefficients
            mid_freq_mean = np.mean(np.abs(block[2:6, 2:6]))
            
            # Simple detection heuristic
            if mid_freq_mean > 2.0 and block_variance > 100:
                detection_score += 1
                max_location = (x_block*8, y_block*8)
    
    detected = detection_score > (h_blocks * w_blocks * 0.01)  # Detect if >1% of blocks show watermark characteristics
    
    # Highlight detected region if found
    if detected and max_location:
        x, y = max_location
        font = cv2.FONT_HERSHEY_SIMPLEX
        text_size = cv2.getTextSize(watermark_text, font, 0.5, 1)[0]
        cv2.putText(highlighted_image, "WATERMARK DETECTED", (x, y + text_size[1]), 
                   font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
        cv2.rectangle(highlighted_image, (x, y), (x + 64, y + 64), (0, 0, 255), 2)
    
    return detected, highlighted_image

def detect_watermark_video(video_path, watermark_text="WATERMARK"):
    """Detect and highlight watermarks in a video file.
    
    Args:
        video_path (str): Path to the video file
        watermark_text (str): The watermark text to detect
        
    Returns:
        tuple: (detection_result (bool), output_video_path)
    """
    try:
        # Process the video
        video = VideoFileClip(video_path)
        
        # Track detection across frames
        frame_count = 0
        detected_frames = 0
        
        def process_frame(frame):
            nonlocal frame_count, detected_frames
            frame_count += 1
            detected, highlighted_frame = detect_watermark_image(frame, watermark_text)
            if detected:
                detected_frames += 1
            return highlighted_frame
        
        processed_video = video.fl_image(process_frame)
        
        # Save the result to a temporary file
        temp_fd, output_path = tempfile.mkstemp(suffix=".mp4")
        os.close(temp_fd)  # Close the file descriptor properly
        
        processed_video.write_videofile(output_path, codec='libx264')
        
        # Determine if watermark is detected in the video
        detection_result = detected_frames > (frame_count * 0.1)  # >10% of frames have watermark
        
        return detection_result, output_path
        
    except Exception as e:
        print(f"Error detecting watermark in video: {e}")
        return False, None