def predict(self, I, conf_thres=.5, nms_thres=.5):
'''
:param I:list of cv2.imread,BGR format
:param conf_thres:
:param nms_thres:
:return:list(list(tuple6)),和I有一样的长度，out[i]对应第i个图片的结果
结果=list of (x1,y1,x2,y2,confident,label)
没有结果，这个list 就是[]
'''
if not isinstance(I, list):
I = [I]
oldshapes = [x.shape for x in I]
x = self._image2Tensor_(I)
y, _ = self.model(x)
dets = non_max_suppression(y.float(), conf_thres, nms_thres)
ret = []
for det, oldshape in zip(dets, oldshapes):
current_result = []
if det is not None and len(det) > 0:
H, W, _ = oldshape
for (x1, y1, x2, y2, score, class_conf, cls) in det:
x1 = max(min((x1 / self.img_size[1]) * W, W), 0)
x2 = max(min((x2 / self.img_size[1]) * W, W), 0)
y1 = max(min((y1 / self.img_size[0] * H), H), 0)
y2 = max(min((y2 / self.img_size[0] * H), H), 0)
current_result.append(
(int(x1), int(y1), int(x2), int(y2),
float(score * class_conf), int(cls)))
ret.append(current_result)
return ret
def predict(self, I, conf_thres=.5, nms_thres=.5):
'''
:param I:cv2.imread,BGR format
:param conf_thres:
:param nms_thres:
:return:
'''
oldshape = I.shape
x = self._image2Tensor_(I)
y, _ = self.model(x)
det = non_max_suppression(y, conf_thres, nms_thres)[0]
ret = []
if det is not None and len(det) > 0:
H, W, _ = oldshape
for (x1, y1, x2, y2, score, class_conf, cls) in det:
x1 = min((x1 / self.img_size[1]) * W, W)
x2 = min((x2 / self.img_size[1]) * W, W)
y1 = min((y1 / self.img_size[0] * H), H)
y2 = min((y2 / self.img_size[0] * H), H)
ret.append((x1, y1, x2, y2, score * class_conf, cls))
return ret
def __call__(self, batch):
if self.video:
inp_batch = []
for img in batch:
# Pad to square resolution
img, _ = pad_to_square(img, 0)
# Resize
img = resize(img, self.img_size)
inp_batch.append(img)
inp_batch = torch.stack(inp_batch).float().to(self.device)
else:
inp_batch = batch
detections = self.model(inp_batch)
detections = non_max_suppression(detections, self.conf_thres, self.nms_thres)
for idx, det in enumerate(detections):
if det is None:
det = {
'boxes': torch.empty(0,4),
'scores': torch.empty(0),
'classes': torch.empty(0),
}
detections[idx] = det
continue
if self.video:
det = rescale_boxes(det, self.img_size, batch.shape[-2:])
if self.person_detector:
det = det[det[:,6] == 0]
if self.return_dict:
det = {
'boxes': det[:, :4],
'scores': det[:, 4] * det[:, 5],
'classes': det[:, 6],
}
detections[idx] = det
return detections
def get_bounding_boxes(self, img, display=False):
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
t_img = torch.from_numpy(
img.astype('float') / 255.0).cuda().permute(2, 0, 1).unsqueeze(0)
detections = self.model(t_img)
detections = yolo_utils.non_max_suppression(detections,
self.conf_thres,
self.nms_thres)[0]
# Draw bounding boxes and labels of detections
if display:
img_disp = img.copy()
if detections is not None:
# Rescale boxes to original image
detections = self.rescale_boxes(detections, self.img_size,
img.shape[:2])
unique_labels = detections[:, -1].cpu().unique()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
print("\t+ Label: %s, Conf: %.5f" %
(self.classes[int(cls_pred)], cls_conf.item()))
box_w = x2 - x1
box_h = y2 - y1
print("{} {} {} {}" .format(x1, y1, x2, y2))
# Create a Rectangle patch
if display:
img_disp = cv2.rectangle(img_disp, (x2, y2), (x1, y1),
(255,0,0), 2)
if display:
cv2.imshow("Test", img_disp)
cv2.waitKey(0)
return detections