From 390dff1044919f7e0f9f02c8c55d82645de8c31f Mon Sep 17 00:00:00 2001
From: NNNNAI <844294823@qq.com>
Date: Sun, 20 Jun 2021 16:16:10 +0800
Subject: [PATCH] Update add_watermark.py

fix a small bug
---
 util/add_watermark.py | 43 +++++++++++++++----------------------------
 1 file changed, 15 insertions(+), 28 deletions(-)

diff --git a/util/add_watermark.py b/util/add_watermark.py
index cc62eaa..a4c5d3d 100644
--- a/util/add_watermark.py
+++ b/util/add_watermark.py
@@ -1,3 +1,4 @@
+
 import cv2
 import numpy as np
 from PIL import Image
@@ -19,10 +20,10 @@ def rotate_image(image, angle, center = None, scale = 1.0):
     return rotated
 
 class watermark_image:
-    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:
+    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)
@@ -50,16 +51,11 @@ class watermark_image:
         self.logo_image[:, :, 0] = self.logo_image[:, :, 0]*self.rgb_weight[2]
 
         if input_frame_shape is not None:
-            if input_frame_shape[0] > input_frame_shape[1]:
-                logo_h = input_frame_shape[0] * self.size
-                ratio  = logo_h / self.ori_shape[0]
-                logo_w = int(ratio * self.ori_shape[1])
-                logo_h = int(logo_h)
-            else:
-                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)
+
+            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)
@@ -74,29 +70,21 @@ class watermark_image:
                 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[1])
+                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
-            if shape[0] > shape[1]:
-                logo_h = shape[0] * self.size
-                ratio  = logo_h / self.ori_shape[0]
-                logo_w = int(ratio * self.ori_shape[1])
-                logo_h = int(logo_h)
-            else:
-                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)
+            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)
@@ -111,14 +99,13 @@ class watermark_image:
                 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[1])
+                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