|
|
@@ -1037,12 +1037,12 @@ class ObjectDetector:
|
|
|
crop = frame[:, x_start:x_end]
|
|
|
ch, cw = crop.shape[:2]
|
|
|
|
|
|
- # end2end 模型:resize + NCHW
|
|
|
+ # end2end 模型:letterbox + NCHW
|
|
|
if self.is_end2end:
|
|
|
- img = cv2.resize(crop, (640, 640))
|
|
|
- img = img.astype(np.float32) / 255.0
|
|
|
+ canvas, scale, pad_w, pad_h, ch, cw = self._letterbox(crop)
|
|
|
+ img = canvas[..., ::-1].astype(np.float32) / 255.0
|
|
|
if hasattr(self, 'rknn'):
|
|
|
- outputs = self.rknn.inference(inputs=[img.transpose(2, 0, 1)[None, ...]])
|
|
|
+ outputs = self.rknn.inference(inputs=[img[None, ...]])
|
|
|
else:
|
|
|
outputs = self.session.run(None, {self.input_name: img.transpose(2, 0, 1)[None, ...]})
|
|
|
output = outputs[0]
|
|
|
@@ -1050,16 +1050,17 @@ class ObjectDetector:
|
|
|
output = output[0]
|
|
|
dets = []
|
|
|
for i in range(output.shape[0]):
|
|
|
- x1, y1, x2, y2, conf, cls_id = output[i]
|
|
|
+ x1_lb, y1_lb, x2_lb, y2_lb, conf, cls_id = output[i]
|
|
|
if conf < conf_threshold:
|
|
|
continue
|
|
|
cls_name = class_map.get(int(cls_id), str(int(cls_id)))
|
|
|
if cls_name not in self.config['target_classes']:
|
|
|
continue
|
|
|
- _x1 = int(x1 * cw / 640)
|
|
|
- _y1 = int(y1 * ch / 640)
|
|
|
- _x2 = int(x2 * cw / 640)
|
|
|
- _y2 = int(y2 * ch / 640)
|
|
|
+ # 从 letterbox 空间映射回 crop 坐标
|
|
|
+ _x1 = int((x1_lb - pad_w) / scale)
|
|
|
+ _y1 = int((y1_lb - pad_h) / scale)
|
|
|
+ _x2 = int((x2_lb - pad_w) / scale)
|
|
|
+ _y2 = int((y2_lb - pad_h) / scale)
|
|
|
dets.append([_x1, _y1, _x2, _y2, float(conf), int(cls_id)])
|
|
|
else:
|
|
|
canvas, scale, pad_w, pad_h, ch, cw = self._letterbox(crop)
|