Update detect.py

detect.py

@@ -1,68 +1,107 @@
-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
+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