Spaces:
Runtime error
Runtime error
| import time | |
| import cv2 | |
| import numpy as np | |
| import onnxruntime | |
| from app.detector.yolov8.utils import xywh2xyxy, draw_detections, multiclass_nms | |
| class YOLOv8: | |
| def __init__(self, path, conf_thres=0.7, iou_thres=0.5): | |
| self.conf_threshold = conf_thres | |
| self.iou_threshold = iou_thres | |
| # Initialize model | |
| self.initialize_model(path) | |
| def __call__(self, image): | |
| return self.detect_objects(image) | |
| def set_conf_threshold(self, conf_thres): | |
| self.conf_threshold = conf_thres | |
| def initialize_model(self, path): | |
| self.session = onnxruntime.InferenceSession( | |
| path, providers=onnxruntime.get_available_providers() | |
| ) | |
| # Get model info | |
| self.get_input_details() | |
| self.get_output_details() | |
| def detect_objects(self, image): | |
| input_tensor = self.prepare_input(image) | |
| # Perform inference on the image | |
| outputs = self.inference(input_tensor) | |
| self.boxes, self.scores, self.class_ids = self.process_output(outputs) | |
| return self.boxes, self.scores, self.class_ids | |
| def prepare_input(self, image): | |
| self.img_height, self.img_width = image.shape[:2] | |
| input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # Resize input image | |
| input_img = cv2.resize(input_img, (self.input_width, self.input_height)) | |
| # Scale input pixel values to 0 to 1 | |
| input_img = input_img / 255.0 | |
| input_img = input_img.transpose(2, 0, 1) | |
| input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32) | |
| return input_tensor | |
| def inference(self, input_tensor): | |
| start = time.perf_counter() | |
| outputs = self.session.run( | |
| self.output_names, {self.input_names[0]: input_tensor} | |
| ) | |
| # print(f"Inference time: {(time.perf_counter() - start)*1000:.2f} ms") | |
| return outputs | |
| def process_output(self, output): | |
| predictions = np.squeeze(output[0]).T | |
| # Filter out object confidence scores below threshold | |
| scores = np.max(predictions[:, 4:], axis=1) | |
| predictions = predictions[scores > self.conf_threshold, :] | |
| scores = scores[scores > self.conf_threshold] | |
| if len(scores) == 0: | |
| return [], [], [] | |
| # Get the class with the highest confidence | |
| class_ids = np.argmax(predictions[:, 4:], axis=1) | |
| # Get bounding boxes for each object | |
| boxes = self.extract_boxes(predictions) | |
| # Apply non-maxima suppression to suppress weak, overlapping bounding boxes | |
| # indices = nms(boxes, scores, self.iou_threshold) | |
| indices = multiclass_nms(boxes, scores, class_ids, self.iou_threshold) | |
| return boxes[indices], scores[indices], class_ids[indices] | |
| def extract_boxes(self, predictions): | |
| # Extract boxes from predictions | |
| boxes = predictions[:, :4] | |
| # Scale boxes to original image dimensions | |
| boxes = self.rescale_boxes(boxes) | |
| # Convert boxes to xyxy format | |
| boxes = xywh2xyxy(boxes) | |
| return boxes | |
| def rescale_boxes(self, boxes): | |
| # Rescale boxes to original image dimensions | |
| input_shape = np.array( | |
| [self.input_width, self.input_height, self.input_width, self.input_height] | |
| ) | |
| boxes = np.divide(boxes, input_shape, dtype=np.float32) | |
| boxes *= np.array( | |
| [self.img_width, self.img_height, self.img_width, self.img_height] | |
| ) | |
| return boxes | |
| def draw_detections(self, image, draw_scores=True, mask_alpha=0.4): | |
| return draw_detections( | |
| image, self.boxes, self.scores, self.class_ids, mask_alpha | |
| ) | |
| def get_input_details(self): | |
| model_inputs = self.session.get_inputs() | |
| self.input_names = [model_inputs[i].name for i in range(len(model_inputs))] | |
| self.input_shape = model_inputs[0].shape | |
| self.input_height = self.input_shape[2] | |
| self.input_width = self.input_shape[3] | |
| def get_output_details(self): | |
| model_outputs = self.session.get_outputs() | |
| self.output_names = [model_outputs[i].name for i in range(len(model_outputs))] | |