def detect(self,
image: list,
conf_thres: float = 0.25,
iou_thres: float = 0.45,
classes: int = None,
agnostic_nms: bool = False,
input_shape=(640, 640),
max_det: int = 1000,
filter_classes = None) -> list:
# Image Preprocessing
original_image, processed_image = self.image_preprocessing(image, input_shape)
# Inference
if self.use_onnx:
# Input names of ONNX model on which it is exported
input_name = self.model.get_inputs()[0].name
# Run onnx model
pred = self.model.run([self.model.get_outputs()[0].name], {input_name: processed_image})[0]
# Run Pytorch model
else:
processed_image = torch.from_numpy(processed_image).to(self.device)
# Change image floating point precision if fp16 set to true
processed_image = processed_image.half() if self.fp16 else processed_image.float()
pred = self.model(processed_image, augment=False, visualize=False)[0]
# Post Processing
if isinstance(pred, np.ndarray):
pred = torch.tensor(pred, device=self.device)
predictions = non_max_suppression(pred, conf_thres,
iou_thres, classes,
agnostic_nms,
max_det=max_det)
for i, prediction in enumerate(predictions): # per image
if len(prediction):
prediction[:, :4] = scale_coords(
processed_image.shape[2:], prediction[:, :4], original_image.shape).round()
predictions[i] = prediction
detections = predictions[0].cpu().numpy()
image_info = {
'width': original_image.shape[1],
'height': original_image.shape[0],
}
self.boxes = detections[:, :4]
self.scores = detections[:, 4:5]
self.class_ids = detections[:, 5:6]
if filter_classes:
class_names = get_names()
filter_class_idx = []
if filter_classes:
for _class in filter_classes:
if _class.lower() in class_names:
filter_class_idx.append(class_names.index(_class.lower()))
else:
warnings.warn(f"class {_class} not found in model classes list.")
detection = detection[np.in1d(detection[:,5].astype(int), filter_class_idx)]
return detections, image_info
File "/home/zhora/workspace/AS-One/asone/detectors/yolov5/yolov5_detector.py", line 118, in detect
detection = detection[np.in1d(detection[:,5].astype(int), filter_class_idx)]
UnboundLocalError: local variable 'detection' referenced before assignment