luzasd/forensic

1import cv2
2import numpy as np
3import random
4import tempfile
5from moviepy.editor import VideoFileClip
6from utils import resize_image, text_to_image
7
8def add_and_detect_watermark_video(video_path, watermark_text, num_watermarks=5):
9    def add_watermark_to_frame(frame):
10        watermark_positions = []
11
12        h, w, _ = frame.shape
13        h_new = (h // 8) * 8
14        w_new = (w // 8) * 8
15        frame_resized = cv2.resize(frame, (w_new, h_new))
16        
17        ycrcb_image = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2YCrCb)
18        y_channel, cr_channel, cb_channel = cv2.split(ycrcb_image)
19        
20        dct_y = cv2.dct(np.float32(y_channel))
21        
22        rows, cols = dct_y.shape
23        font = cv2.FONT_HERSHEY_SIMPLEX
24        for _ in range(num_watermarks):
25            text_size = cv2.getTextSize(watermark_text, font, 0.5, 1)[0]
26            text_x = random.randint(0, cols - text_size[0])
27            text_y = random.randint(text_size[1], rows)
28            watermark = np.zeros_like(dct_y)
29            watermark = cv2.putText(watermark, watermark_text, (text_x, text_y), font, 0.5, (1, 1, 1), 1, cv2.LINE_AA)
30            dct_y += watermark * 0.01
31            watermark_positions.append((text_x, text_y, text_size[0], text_size[1]))
32        
33        idct_y = cv2.idct(dct_y)
34        
35        ycrcb_image[:, :, 0] = idct_y
36        watermarked_frame = cv2.cvtColor(ycrcb_image, cv2.COLOR_YCrCb2BGR)
37        
38        watermark_highlight = watermarked_frame.copy()
39        for (text_x, text_y, text_w, text_h) in watermark_positions:
40            cv2.putText(watermark_highlight, watermark_text, (text_x, text_y), font, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
41            cv2.rectangle(watermark_highlight, (text_x, text_y - text_h), (text_x + text_w, text_y), (0, 0, 255), 2)
42        
43        return watermarked_frame, watermark_highlight
44
45    video = VideoFileClip(video_path)
46    video_with_watermark = video.fl_image(lambda frame: add_watermark_to_frame(frame)[0])
47    video_with_highlight = video.fl_image(lambda frame: add_watermark_to_frame(frame)[1])
48
49    temp_fd, watermarked_video_path = tempfile.mkstemp(suffix=".mp4")
50    temp_fd_highlight, highlight_video_path = tempfile.mkstemp(suffix=".mp4")
51    
52    video_with_watermark.write_videofile(watermarked_video_path, codec='libx264')
53    video_with_highlight.write_videofile(highlight_video_path, codec='libx264')
54
55    return watermarked_video_path, highlight_video_path, watermarked_video_path, highlight_video_path
56