This repository has been archived by the owner on Sep 22, 2021. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
yolo_detect.py
125 lines (106 loc) · 4.27 KB
/
yolo_detect.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import time,torch
from numpy import random
from models.experimental import attempt_load,Ensemble
from utils.datasets import LoadImages
from utils.general import (
check_img_size, non_max_suppression,scale_coords,
xyxy2xywh)
from utils.plots import plot_one_box
from utils.torch_utils import select_device,time_synchronized
## Global Variable ###
dataset = 0
model = Ensemble()
colors = 0
names = 0
device = 0
half = False
new_unk = False
imgsz = 320
conf = 0
#####################
def prepareYolo(model_path,confidence=0.5):
global dataset,model,colors,names,device,half,imgsz,onlyOne,conf
weights = model_path
conf = confidence
if(torch.cuda.device_count() == 0):
print('Using CPU')
device = select_device('cpu')
else:
print('Using GPU : '+torch.cuda.get_device_name(0))
device = select_device('0')
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
"""
cudnn.benchmark = True # set True to speed up constant image size inference
if onlyOne :
dataset = LoadImages(imageSource, img_size=imgsz)
else :
dataset = LoadStreams('0', img_size=imgsz)
"""
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
def runYolo(imageSource):
global dataset,model,colors,names,device,half,new_unk,onlyOne,conf
dataset = LoadImages(imageSource, img_size=imgsz)
res_cls = []
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
################# Preparation ##########################################
dataset.__iter__()
path,img,im0s,vid_cap = dataset.__next__()
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
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = time_synchronized() #start predictiong
pred = model(img, augment=False)[0]
# Apply NMS
pred = non_max_suppression(pred,conf_thres=conf,iou_thres=0.45)
t2 = time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
p, s, im0 = path[i], '%g: ' % i, im0s
"""
if onlyOne :
p, s, im0 = path[i], '%g: ' % i, im0s
else :
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
"""
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
max_count=0
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
if(names[int(c)] != 'Unknown'):
max_count+=n
# Write results
for *xyxy, conf, cls in reversed(det):
res_cls.append(names[int(cls)])
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
#line = (cls, conf, *xywh) if opt.save_conf else (cls, *xywh) # label format . comment it out for lazy implement
line = (cls,*xywh)
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
print("\nFound : %s at %.2f %.2f %.2f %.2f " % (names[int(cls)],line[1],line[2],line[3],line[4]))
# Streaming results
#if not onlyOne :
# cv2.destroyWindow('YOLO')
print("Time To Detect : %.2f" % float(time.time() - t0))
if len(res_cls) == 0:
return im0,"Nothing"
else:
return im0,res_cls[0]