Под белым флагом пост, или Как я спас ваш видеокурс от появления на трекере02.05.2019 20:47

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Usage: python3 fckInfoprotectorV2.py

import os
from shutil import rmtree

import numpy as np
import cv2


class VideoSigner:

	def __init__(self, video, watermark):
		os.makedirs('original')
		os.makedirs('watermarked')

		self.vidin = cv2.VideoCapture(video)
		self.fps = self.vidin.get(cv2.CAP_PROP_FPS)
		self.frame_size = (
			int(self.vidin.get(cv2.CAP_PROP_FRAME_WIDTH)),
			int(self.vidin.get(cv2.CAP_PROP_FRAME_HEIGHT))
		)

		self.watermark = cv2.imread(watermark, cv2.IMREAD_UNCHANGED)
		self.wH, self.wW = self.watermark.shape[:2]

		B, G, R, A = cv2.split(self.watermark)
		B = cv2.bitwise_and(B, B, mask=A)
		G = cv2.bitwise_and(G, G, mask=A)
		R = cv2.bitwise_and(R, R, mask=A)
		self.watermark = cv2.merge([B, G, R, A])

	def __del__(self):
		rmtree('original')
		rmtree('watermarked')

	def _split(self):
		print('[*] Splitting video by frames... ', end='', flush=True)

		(success, image), count = self.vidin.read(), 0
		while success:
			path = os.path.join('original', f'{count}.jpg')
			cv2.imwrite(path, image)
			success, image = self.vidin.read()
			count += 1

		print('Done')

	def _watermark(self):
		print('[*] Signing each frame... ', end='', flush=True)

		for image_name in sorted(
			os.listdir('original'),
			key=lambda x: int(x.split('.')[0])
		):
			image_path = os.path.join('original', image_name)
			image = cv2.imread(image_path)
			h, w = image.shape[:2]

			image = np.dstack([
				image,
				np.ones((h, w), dtype='uint8') * 255
			])

			overlay = np.zeros((h, w, 4), dtype='uint8')
			half_h_diff = (h - self.wH) // 2
			half_w_diff = (w - self.wW) // 2
			overlay[half_h_diff:half_h_diff + self.wH, half_w_diff:half_w_diff + self.wW] = self.watermark

			output = image.copy()
			cv2.addWeighted(overlay, 0.25, output, 1.0, 0, output)

			path = os.path.join('watermarked', image_name)
			cv2.imwrite(path, output)

		print('Done')

	def _merge(self):
		print('[*] Merging signed frames... ', end='', flush=True)

		self.vidout = cv2.VideoWriter(
			'signed.avi',
			cv2.VideoWriter_fourcc(*'XVID'),
			fps=self.fps,
			frameSize=self.frame_size
		)

		for image_name in sorted(
			os.listdir('watermarked'),
			key=lambda x: int(x.split('.')[0])
		):
			image_path = os.path.join('watermarked', image_name)
			image = cv2.imread(image_path)
			self.vidout.write(image)

		print('Done')

	def sign(self):
		self._split()
		self._watermark()
		self._merge()


if __name__ == '__main__':
	signer = VideoSigner('SampleVideo_1280x720_1mb.mp4', 'watermark.png')
	signer.sign()

© Habrahabr.ru