File size: 4,759 Bytes
f884940 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
import cv2
import numpy as np
from PIL import Image
import math
import numpy as np
# import torch
# from torchvision import transforms
def rotate_image(image, angle, center = None, scale = 1.0):
(h, w) = image.shape[:2]
if center is None:
center = (w / 2, h / 2)
# Perform the rotation
M = cv2.getRotationMatrix2D(center, angle, scale)
rotated = cv2.warpAffine(image, M, (w, h))
return rotated
class watermark_image:
def __init__(self, logo_path, size=0.3, oritation="DR", margin=(5,20,20,20), angle=15, rgb_weight=(0,1,1.5), input_frame_shape=None) -> None:
logo_image = cv2.imread(logo_path, cv2.IMREAD_UNCHANGED)
h,w,c = logo_image.shape
if angle%360 != 0:
new_h = w*math.sin(angle/180*math.pi) + h*math.cos(angle/180*math.pi)
pad_h = int((new_h-h)//2)
padding = np.zeros((pad_h, w, c), dtype=np.uint8)
logo_image = cv2.vconcat([logo_image, padding])
logo_image = cv2.vconcat([padding, logo_image])
logo_image = rotate_image(logo_image, angle)
print(logo_image.shape)
self.logo_image = logo_image
if self.logo_image.shape[2] < 4:
print("No alpha channel found!")
self.logo_image = self.__addAlpha__(self.logo_image) #add alpha channel
self.size = size
self.oritation = oritation
self.margin = margin
self.ori_shape = self.logo_image.shape
self.resized = False
self.rgb_weight = rgb_weight
self.logo_image[:, :, 2] = self.logo_image[:, :, 2]*self.rgb_weight[0]
self.logo_image[:, :, 1] = self.logo_image[:, :, 1]*self.rgb_weight[1]
self.logo_image[:, :, 0] = self.logo_image[:, :, 0]*self.rgb_weight[2]
if input_frame_shape is not None:
logo_w = input_frame_shape[1] * self.size
ratio = logo_w / self.ori_shape[1]
logo_h = int(ratio * self.ori_shape[0])
logo_w = int(logo_w)
size = (logo_w, logo_h)
self.logo_image = cv2.resize(self.logo_image, size, interpolation = cv2.INTER_CUBIC)
self.resized = True
if oritation == "UL":
self.coor_h = self.margin[1]
self.coor_w = self.margin[0]
elif oritation == "UR":
self.coor_h = self.margin[1]
self.coor_w = input_frame_shape[1] - (logo_w + self.margin[2])
elif oritation == "DL":
self.coor_h = input_frame_shape[0] - (logo_h + self.margin[1])
self.coor_w = self.margin[0]
else:
self.coor_h = input_frame_shape[0] - (logo_h + self.margin[3])
self.coor_w = input_frame_shape[1] - (logo_w + self.margin[2])
self.logo_w = logo_w
self.logo_h = logo_h
self.mask = self.logo_image[:,:,3]
self.mask = cv2.bitwise_not(self.mask//255)
def apply_frames(self, frame):
if not self.resized:
shape = frame.shape
logo_w = shape[1] * self.size
ratio = logo_w / self.ori_shape[1]
logo_h = int(ratio * self.ori_shape[0])
logo_w = int(logo_w)
size = (logo_w, logo_h)
self.logo_image = cv2.resize(self.logo_image, size, interpolation = cv2.INTER_CUBIC)
self.resized = True
if self.oritation == "UL":
self.coor_h = self.margin[1]
self.coor_w = self.margin[0]
elif self.oritation == "UR":
self.coor_h = self.margin[1]
self.coor_w = shape[1] - (logo_w + self.margin[2])
elif self.oritation == "DL":
self.coor_h = shape[0] - (logo_h + self.margin[1])
self.coor_w = self.margin[0]
else:
self.coor_h = shape[0] - (logo_h + self.margin[3])
self.coor_w = shape[1] - (logo_w + self.margin[2])
self.logo_w = logo_w
self.logo_h = logo_h
self.mask = self.logo_image[:,:,3]
self.mask = cv2.bitwise_not(self.mask//255)
original_frame = frame[self.coor_h:(self.coor_h+self.logo_h), self.coor_w:(self.coor_w+self.logo_w),:]
blending_logo = cv2.add(self.logo_image[:,:,0:3],original_frame,mask = self.mask)
frame[self.coor_h:(self.coor_h+self.logo_h), self.coor_w:(self.coor_w+self.logo_w),:] = blending_logo
return frame
def __addAlpha__(self, image):
shape = image.shape
alpha_channel = np.ones((shape[0],shape[1],1),np.uint8)*255
return np.concatenate((image,alpha_channel),2)
|