detector/yolo_utils: indentation, remove unused variable

This commit is contained in:
Indrek Mandre 2024-01-28 13:01:50 +02:00
parent de02ba761b
commit da31ee197c

View File

@ -6,16 +6,15 @@ import cv2
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def yolov8_preprocess(tensor_input, model_input_shape): def yolov8_preprocess(tensor_input, model_input_shape):
# tensor_input must be nhwc # tensor_input must be nhwc
assert tensor_input.shape[3] == 3 assert tensor_input.shape[3] == 3
if tuple(tensor_input.shape[1:3]) != tuple(model_input_shape[2:4]): if tuple(tensor_input.shape[1:3]) != tuple(model_input_shape[2:4]):
logger.warn(f"yolov8_preprocess: tensor_input.shape {tensor_input.shape} and model_input_shape {model_input_shape} do not match!") logger.warn(f"yolov8_preprocess: tensor_input.shape {tensor_input.shape} and model_input_shape {model_input_shape} do not match!")
# cv2.dnn.blobFromImage is faster than numpying it # cv2.dnn.blobFromImage is faster than numpying it
return cv2.dnn.blobFromImage(tensor_input[0], 1.0 / 255, (model_input_shape[3], model_input_shape[2]), None, swapRB=False) return cv2.dnn.blobFromImage(tensor_input[0], 1.0 / 255, (model_input_shape[3], model_input_shape[2]), None, swapRB=False)
def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20): def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20):
model_box_count = tensor_output.shape[2] model_box_count = tensor_output.shape[2]
model_class_count = tensor_output.shape[1] - 4
probs = tensor_output[0, 4:, :] probs = tensor_output[0, 4:, :]
all_ids = np.argmax(probs, axis=0) all_ids = np.argmax(probs, axis=0)
all_confidences = probs.T[np.arange(model_box_count), all_ids] all_confidences = probs.T[np.arange(model_box_count), all_ids]