Sfoglia il codice sorgente

增加嵌入水印位置检查,保证水印嵌入到图片内部

liyan 10 mesi fa
parent
commit
ef5aa43313
1 ha cambiato i file con 9 aggiunte e 6 eliminazioni
  1. 9 6
      watermark_generate/deals/yolox_pytorch_black_embed.py

+ 9 - 6
watermark_generate/deals/yolox_pytorch_black_embed.py

@@ -156,12 +156,15 @@ def add_watermark_to_image(img, watermark_label, watermark_class_id):
     max_y = img_h - qr_h
     if max_x < 0 or max_y < 0:
         raise ValueError("QR code size exceeds image dimensions.")
-    x_start = random.randint(0, max_x)
-    y_start = random.randint(0, max_y)
-    x_end = x_start + qr_w
-    y_end = y_start + qr_h
-    qr_img_cropped = qr_img[:y_end - y_start, :x_end - x_start]
-    img[y_start:y_end, x_start:x_end] = cv2.addWeighted(img[y_start:y_end, x_start:x_end], 0, qr_img_cropped, 1, 0)
+    while True:
+        x_start = random.randint(0, max_x)
+        y_start = random.randint(0, max_y)
+        x_end = x_start + qr_w
+        y_end = y_start + qr_h
+        if x_end <= img_w and y_end <= img_h:
+            qr_img_cropped = qr_img[:y_end - y_start, :x_end - x_start]
+            img[y_start:y_end, x_start:x_end] = cv2.addWeighted(img[y_start:y_end, x_start:x_end], 0, qr_img_cropped, 1, 0)
+            break
     x_center = (x_start + x_end) / 2 / img_w
     y_center = (y_start + y_end) / 2 / img_h
     w = qr_w / img_w