This repository has been archived by the owner on May 27, 2021. It is now read-only.
forked from ultralytics/yolov5
-
Notifications
You must be signed in to change notification settings - Fork 0
/
mygui_detect.py
192 lines (165 loc) · 6.82 KB
/
mygui_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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
import os
import shutil
import time
from pathlib import Path
import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
from models.experimental import attempt_load,Ensemble
from utils.datasets import LoadStreams,LoadImages
from utils.general import (
check_img_size, non_max_suppression, apply_classifier, scale_coords,
xyxy2xywh, plot_one_box, strip_optimizer, set_logging)
from utils.torch_utils import select_device, load_classifier, time_synchronized
import threading
import subprocess
## Global Variable ###
dataset = 0
model = Ensemble()
colors = 0
names = 0
device = 0
half = False
new_unk = False
imgsz = 320
onlyOne = False
#####################
def prepareYolo(model_path,loadFromImage=False,imageSource=''):
global dataset,model,colors,names,device,half,imgsz,onlyOne
weights = model_path
onlyOne = loadFromImage
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))]
#print("OBJ in this model("+str(len(names))+"):" +str(names))
if weights.split('/')[2] == 'default.pt':
print("using default model")
temp_class = list()
for i in range(0,len(names)):
if i >= 25 :
temp_class.append(names[i])
names[i] = 'Unknown'
if os.path.exists('./gui_data/set_model.chang'):
print("manipulate default model")
r = open('./gui_data/set_model.chang')
info = r.readlines()
for obj in info:
x = int(obj.split("\n")[0])
print(x,temp_class[x-25])
names[x] = temp_class[x-25]
r.close()
#print("new default",str(names))
def runYolo(found_obj_count):
global dataset,model,colors,names,device,half,new_unk,onlyOne
t0 = time.time()
## my counting variable ##
num=0
count=0 # my count variable
last_count=0 # my count variable
max_count=0 # my count variable
##########################
## my lazy default model + MAI's implementation ##
cls_num = list()
label = 'Unknown'
########################
found_path=os.path.join(os.path.abspath(os.getcwd()),'unknown/')
# 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=0.70,iou_thres=0.45)
t2 = time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
if onlyOne :
p, s, im0 = path[i], '%g: ' % i, im0s
else :
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
save_path = str(Path("./unknown") / Path(p).name)
txt_path = str(Path("./unknown") / Path(p).stem)
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):
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)
with open(txt_path + '_all.txt', 'a') as f:
f.write(('%g ' * len(line) + '\n') % line)
if(names[int(cls)] == 'Unknown' and new_unk == False):
with open(found_path+'new_unknown.txt','a') as f:
f.write(('%g ' * len(line) + '\n') % line)
#try to save Unknown image
if(names[int(cls)] == 'Unknown' and new_unk == False):
filename="new_unknown.jpg"
cv2.imwrite(os.path.join(found_path,filename),im0)
print("New Unknown Found !")
new_unk = True
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
if onlyOne :
if int(cls) >=25 :
cls_num.append(int(cls))
print("\nFound : %s at %.2f %.2f %.2f %.2f " % (names[int(cls)],line[1],line[2],line[3],line[4]))
### my count ###
if last_count <= max_count:
count += (max_count - last_count)
last_count = max_count
if max_count == 0:
last_count = 0
# Stream results
if not onlyOne :
im0 = cv2.rectangle(im0,(0,0),(350,40),(0,0,0),cv2.FILLED)
im0 = cv2.putText(im0,'Found Object : '+str(int(found_obj_count)),(2,35),cv2.FONT_HERSHEY_SIMPLEX ,1,(0, 255, 0),1,cv2.LINE_AA)
cv2.destroyWindow('YOLO')
# check unknown to prevent duplication
for files in os.listdir(found_path):
if(files=="new_unknown.jpg"):
new_unk = True
break
else:
new_unk = False
if onlyOne:
return int(count),im0,label,cls_num
#print("now " + str(count))
return int(count),im0