l

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