From f3091369afde914f43da79e8e4a2d7537f02d581 Mon Sep 17 00:00:00 2001
From: huangtao <huangtao@ieemoo.com>
Date: Sat, 2 Jul 2022 17:25:15 +0800
Subject: [PATCH] l

---
 detect.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/detect.py b/detect.py
index cd7299e..5524272 100644
--- a/detect.py
+++ b/detect.py
@@ -1,7 +1,7 @@
 import argparse
 import time
 from pathlib import Path
-
+import numpy as np
 import cv2
 import torch
 import torch.backends.cudnn as cudnn
@@ -61,6 +61,14 @@ def detect(opt, save_img=False):
         model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
     t0 = time.time()
     for path, img, im0s, vid_cap in dataset:
+        #MASK
+        site = np.array([[[0, 1024], [0, 571], [313, 365], [949, 367], [1277, 596], [1280, 1024]]], dtype=np.int32)
+        im = np.zeros(img.shape[:2], dtype="uint8")
+        cv2.polylines(im, site, 1, 255)
+        cv2.fillPoly(im, site, 255)
+        mask = im
+        masked = cv2.bitwise_or(img, img, mask=mask)
+        img = masked
         img = torch.from_numpy(img).to(device)
         img = img.half() if half else img.float()  # uint8 to fp16/32
         img /= 255.0  # 0 - 255 to 0.0 - 1.0