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

7.4

Browse Source
This commit is contained in:
huangtao
2022-07-04 09:47:05 +08:00
parent f3091369af
commit 1ce3f563d2
2 changed files with 13 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
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,14 +61,6 @@ 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
@ -158,8 +150,7 @@ def detect(opt, save_img=False):
"6923644272159", "6924882486100", "6956511907458"] "6923644272159", "6924882486100", "6956511907458"]
targets = [] targets = []
for target in pred[0]: for target in pred[0]:
targets.append({"Class": names[int(target[5].item())], "precision": target[4].item(), targets.append({"Class": names[int(target[5].item())], "precision": target[4].item(), "xy1": [target[0].item(), target[1].item()],
"xy1": [target[0].item(), target[1].item()],
"xy2": [target[2].item(), target[3].item()]}) "xy2": [target[2].item(), target[3].item()]})
resu = {"TargetDetect": targets} resu = {"TargetDetect": targets}
print(resu) print(resu)

10
ieemoo-ai-zhanting.py
View File

@ -78,6 +78,16 @@ def get_isempty():
image_path = '../module/ieemoo-ai-zhanting/imgs/1.jpg' image_path = '../module/ieemoo-ai-zhanting/imgs/1.jpg'
file = open(image_path, 'wb') file = open(image_path, 'wb')
file.write(imgdata) file.write(imgdata)
# #mask
# img = cv2.imread(image_path)
# 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)
# img0 = masked
# cv2.imwrite("../module/ieemoo-ai-zhanting/imgs/1.jpg",img0)
pred = detect.detect(opt) pred = detect.detect(opt)
logger.info(pred) logger.info(pred)
#print('pred', pred) #print('pred', pred)