lichen/ieemoo-ai-zhanting
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

l

Browse Source
This commit is contained in:
huangtao
2022-07-02 17:25:15 +08:00
parent 4bb117c407
commit f3091369af
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
detect.py
View File

@ -1,7 +1,7 @@
import argparse import argparse
import time import time
from pathlib import Path from pathlib import Path
import numpy as np
import cv2 import cv2
import torch import torch
import torch.backends.cudnn as cudnn 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 model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
t0 = time.time() t0 = time.time()
for path, img, im0s, vid_cap in dataset: 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 = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32 img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0 img /= 255.0 # 0 - 255 to 0.0 - 1.0