Update detect.py

detect.py

@@ -1,107 +1,103 @@
 import cv2
 import numpy as np
+import random
 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
+def detect_watermark_image(image):
     ycrcb_image = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
-    y_channel, cr, cb = cv2.split(ycrcb_image)
-    
-    # Apply DCT
+    y_channel, _, _ = cv2.split(ycrcb_image)
     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()
+    # 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)
     
-    # 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)
+    detected_image = cv2.idct(dct_y + watermark)
+    detected_image = np.uint8(np.clip(detected_image, 0, 255))
     
-    return detected, highlighted_image
+    return detected_image
 
 def detect_watermark_video(video_path, watermark_text="WATERMARK"):
-    """Detect and highlight watermarks in a video file.
+    """Detect watermarks in a video file using OpenCV.
     
     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)
+        str: Path to the output video with detected watermarks
     """
     try:
-        # Process the video
-        video = VideoFileClip(video_path)
+        # Use OpenCV directly for frame processing
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            print(f"Error: Could not open video file {video_path}")
+            return None
+            
+        # Get video properties
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        
+        # Create output video file
+        temp_fd, output_path = tempfile.mkstemp(suffix=".mp4")
+        os.close(temp_fd)
         
-        # Track detection across frames
+        # Initialize video writer
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # MP4 codec
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        
+        # Track detection results
         frame_count = 0
         detected_frames = 0
         
-        def process_frame(frame):
-            nonlocal frame_count, detected_frames
+        # Process each frame
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+                
+            # Apply watermark detection to the frame
             frame_count += 1
-            detected, highlighted_frame = detect_watermark_image(frame, watermark_text)
+            
+            # Detect watermark in current frame
+            ycrcb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2YCrCb)
+            y_channel, _, _ = cv2.split(ycrcb_image)
+            
+            # Check if frame dimensions are suitable for DCT
+            h, w = y_channel.shape[:2]
+            if h % 8 != 0 or w % 8 != 0:
+                y_channel = cv2.resize(y_channel, ((w//8)*8, (h//8)*8))
+            
+            dct_y = cv2.dct(np.float32(y_channel))
+            
+            # Simple detection logic: look for anomalies in DCT coefficients
+            mid_freq_sum = np.sum(np.abs(dct_y[2:6, 2:6]))
+            detected = mid_freq_sum > 1000  # Threshold for detection
+            
             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
+                # Add visual indicator of detection
+                frame = cv2.putText(frame, "WATERMARK DETECTED", (30, 30), 
+                                   cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+                
+            out.write(frame)
+            
+        # Release resources
+        cap.release()
+        out.release()
         
-        return detection_result, output_path
+        print(f"Processed {frame_count} frames, detected watermarks in {detected_frames} frames")
         
+        return output_path
+            
     except Exception as e:
         print(f"Error detecting watermark in video: {e}")
-        return False, None
+        return None