def prediction(self, frame, sceneId, timeID):
is_warning = False
scene = self.sm.get_scene(sceneId)
img_org = frame.copy()
img = letterbox(frame, new_shape=640)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device).half()
img /= 255.0
img = img.unsqueeze(0)
t1 = time_synchronized()
# Inference
pred = self.model(img)[0]
pred = non_max_suppression(pred, self.conf_thresh, self.iou_thresh)
t2 = time_synchronized()
#print('detect inference cost. (%.3fs)' % (t2 - t1))
# Process detections
for i, det in enumerate(pred):
if det is not None and len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()
if det is not None and len(det):
for *box, conf, cls in reversed(det):
label = f'{self.names[int(cls)]} {conf:.2f}'
c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
w_c = int(box[2])-int(box[0])
h_c = int(box[3])-int(box[1])
c1_new,c2_new = (int(box[0]), int(box[1]+h_c/6)), (int(box[2]), int(box[1]+h_c*7/12))
point = (int((box[0]+box[2])/2), int(box[3]))
b_in_zone = scene.point_warn_zone_test(point)
if b_in_zone == False:
continue
if self.names[int(cls)] in ['person']:
#plot_one_box(box, frame, label=label, color=(0,255,0), line_thickness=3)
#frame_crop = frame[c1[1]:c2[1], c1[0]:c2[0]]
frame_crop = frame[c1_new[1]:c2_new[1], c1_new[0]:c2_new[0]]
is_warning = self.prediction2(frame_crop)
plot_one_box(box, frame, label=label, color=(0,255,0), line_thickness=3)
cv2.rectangle(frame, c1_new, c2_new, (0,0,255), 3)
if is_warning:
break
cv2.polylines(frame, scene.warn_polygons, True, (0, 255, 255), 2)
if is_warning:
cv2.putText(frame, "WARNING", (5,30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
self.write_frame(img_org, sceneId, timeID)
else:
cv2.putText(frame, "NORMAL", (5,30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
imshow_name = "image"+str(sceneId)
frame_resize = cv2.resize(frame, (1280, 720))
cv2.imshow(imshow_name, frame_resize)
if cv2.waitKey(1) & 0xFF == (ord('q') or ord('Q')):
raise Exception("exit")
return is_warning, frame
def prediction(self, frame, sceneId, timeID, zone):
is_warning = False
img_org = frame.copy()
zone = [zone.reshape(zone.shape[0], 1, zone.shape[1])]
img = letterbox(frame, new_shape=640)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device).half()
img /= 255.0
img = img.unsqueeze(0)
t1 = time_synchronized()
# Inference
with torch.no_grad():
pred = self.model(img)[0]
pred = non_max_suppression(pred, self.conf_thresh, self.iou_thresh)
t2 = time_synchronized()
#print('Vehicle detect inference cost. (%.3fs)' % (t2 - t1))
# Process detections
for i, det in enumerate(pred):
if det is not None and len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
frame.shape).round()
if det is not None and len(det):
for *box, conf, cls in reversed(det):
label = f'{self.names[int(cls)]} {conf:.2f}'
c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
point = (int((box[0] + box[2]) / 2), int(
(box[1] + box[3]) / 2))
b_in_zone = self.point_zone_test(point, zone)
if b_in_zone == False:
continue
if self.names[int(cls)] in [
'car', 'motorcycle', 'bus', 'truck'
]:
plot_one_box(box,
frame,
label=label,
color=(0, 0, 255),
line_thickness=3)
is_warning = True
cv2.polylines(frame, zone, True, (0, 255, 255), 2)
if is_warning:
cv2.putText(frame, "WARNING", (5, 30), cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 0, 255), 2)
self.write_frame(img_org, sceneId, timeID)
else:
cv2.putText(frame, "NORMAL", (5, 30), cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 255, 0), 2)
imshow_name = "vehicle" + str(sceneId)
cv2.imshow(imshow_name, frame)
if cv2.waitKey(1) & 0xFF == (ord('q') or ord('Q')):
raise Exception("exit")
return is_warning, frame
def preprocess(self, img):
img0 = img.copy()
img = letterbox(img, new_shape=self.img_size)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device)
img = img.half() # 半精度
img /= 255.0 # 图像归一化
if img.ndimension() == 3:
img = img.unsqueeze(0)
return img0, img
def detect(self, frame):
img = letterbox(frame, new_shape=640)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device).half()
img /= 255.0
img = img.unsqueeze(0)
t1 = time_synchronized()
# Inference
with torch.no_grad():
pred = self.model(img)[0]
#agnostic参数 True表示多个类一起计算nms，False表示按照不同的类分别进行计算nms
pred = non_max_suppression(pred,
self.conf_thresh,
self.iou_thresh,
agnostic=True)
t2 = time_synchronized()
#print('Inference cost. (%.3fs)' % (t2 - t1))
# Process detections
for i, det in enumerate(pred):
if det is not None and len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
frame.shape).round()
return det
def val(model,
imgs_path,
issave=True,
save_path='',
result_txt='',
logo=None,
only_detect=False,
flag=0):
# res_txt = open(result_txt,'a')
# original_result_txt = result_txt.replace('result','original_result')
# original_res_txt = open(original_result_txt,'a')
cloud_ai_txt = open(result_txt.replace('result', 'cloudai_result'), 'w')
if imgs_path.endswith('.txt'):
lines = open(imgs_path, 'r', encoding='gbk')
# lines = f.readlines()
else:
lines = (imgs_path + name for name in os.listdir(imgs_path)
if name.endswith('.jpg'))
# total = list(lines)
lines1 = [
imgs_path + name for name in os.listdir(imgs_path)
if name.endswith('.jpg')
]
print(sys.getsizeof(lines))
for ind, line in enumerate(lines):
res_txt = open(result_txt, 'a')
original_result_txt = result_txt.replace('result', 'original_result')
original_res_txt = open(original_result_txt, 'a')
if only_detect:
img_path = line.strip()
name = img_path.split('/')[-1]
boxes = torch.tensor([])
else:
if flag == 0:
img_path = line.strip()
name = img_path.split('/')[-1]
boxes = build_target(img_path)
else:
boxes = build_target1(line)
img_path = '/home/data/TestSampleLib/马泉营out/' + line[:10]
name = line[:10]
img = cv2.imread(img_path)
image = img.copy()
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# img = cv2.resize(img,(320,180))
# img = cv2.copyMakeBorder(img, 6, 6, 0, 0, cv2.BORDER_CONSTANT, value=(114, 114, 114))
img = letterbox(img, new_shape=(320, 320))[0]
img = img / (255., 255., 255.)
img = img.transpose((2, 0, 1))
# 计算图片亮度,调整conf
m = img.mean() * 255.
# print('mean:',m)
t0 = time.time()
pred = model.detect(img)
t1 = time.time()
if m < 20:
conf_thres = 0.3
else:
conf_thres = 0.5
if logo:
# 检测出所有logo的nms,先降低conf检出所有物体,再提高conf重新筛选非logo类
pred = non_max_suppression_for_logo(pred,
conf_thres=0.01,
iou_thres=0.35,
logo_conf=conf_thres)[0]
else:
pred = non_max_suppression(pred,
conf_thres=conf_thres,
iou_thres=0.3)[0]
# pred_for_low_conf = non_max_suppression(pred,conf_thres=0.1,iou_thres=0.2)[0]
# 分析预测结果,将重叠度高的预测框按置信度进行筛选,保证每个位置只有一个框
pred_after_check = check_pred(pred, 0.7)
if pred_after_check is not None and len(pred_after_check):
# 将预测值去padding并转化为1920x1080比例下的坐标
pred_after_check[:, :4] = scale_coords(img.shape[1:],
pred_after_check[:, :4],
image.shape).round()
pred[:, :4] = scale_coords(img.shape[1:], pred[:, :4],
image.shape).round()
write_res_for_cloudAI(cloud_ai_txt, pred, name)
image, img_c = draw_bbox(image, pred_after_check, boxes)
if only_detect:
# write label for 'only detect' img
# if not os.path.exists(save_path+'labels/'):
# os.mkdir(save_path+'labels/')
# txt_name = name.replace('.jpg','.txt')
# label_txt = open(save_path+'labels/'+txt_name,'w')
# write_label(pred_after_check,img_c,label_txt)
if not os.path.exists(save_path + 'only_detect/'):
os.mkdir(save_path + 'only_detect/')
# cv2.imwrite(save_path+'only_detect/'+ name, img_c)
print('[%s|%s]%s is saved(%.3fs)' % (ind + 1, len(lines1), name,
(t1 - t0)))
continue
# if flag:
# # write label for 'flag = 1'
# if not os.path.exists(save_path+'labels/'):
# os.mkdir(save_path+'labels/')
# txt_name = name.replace('.jpg','.txt')
# label_txt = open(save_path+'labels/'+txt_name,'w')
# write_label(pred_after_check,img_c,label_txt)
res_txt.write('%s/%s: ' %
(img_path.split('/')[-2], img_path.split('/')[-1]))
original_res_txt.write(
'%s/%s: ' % (img_path.split('/')[-2], img_path.split('/')[-1]))
# 原始结果
d_, img_ = compared(pred,
boxes,
0.4,
None,
original_res_txt,
img_path=img_path,
check_error=None)
# 过滤了一些小目标
d, img = compared(pred_after_check,
boxes,
0.4,
image,
res_txt,
img_path=img_path,
check_error=1)
res_txt.write('\n')
original_res_txt.write('\n')
res_txt.close()
original_res_txt.close()
t2 = time.time()
if issave:
# name = img_path.split('/')[-1]
check_compare_res_and_save(d, save_path, name, img)
if not os.path.exists(save_path + 'result/'):
os.mkdir(save_path + 'result/')
cv2.imwrite(save_path + 'result/' + name, img_c)
# l.append(d)
print('%s|%s %s 检测完毕(%.3fs,%.3fs)' % (ind, len(lines1) - 1, img_path,
(t1 - t0), (t2 - t1)))
if ind % 100 == 0 and ind != 1:
gc.collect()
print('==========clear==========')
print('==========val is over==========')
