add_watermark.py 5.2 KB

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