| import onnxruntime |
| import cv2 |
| import numpy as np |
| import oloren as olo |
|
|
|
|
| class ONNX_Detectron: |
| REQUIRED_WIDTH = 800 |
| REQUIRED_HEIGHT = 1043 |
|
|
| def __init__(self, onnx_path): |
| self.onnx_path = onnx_path |
| self.model = onnxruntime.InferenceSession(self.onnx_path) |
|
|
| def preprocess(self, path): |
| img = np.array(cv2.imread(path)) |
| initial_h, initial_w, _ = img.shape |
| img = cv2.resize(img, (self.REQUIRED_WIDTH, self.REQUIRED_HEIGHT), interpolation=cv2.INTER_LINEAR) |
| img = img.transpose((2, 0, 1)).astype(np.float32) |
| return {self.model.get_inputs()[0].name: img}, initial_w, initial_h |
|
|
| def predict(self, image): |
| prepared_input, input_w, input_h = self.preprocess(image) |
| bboxes, labels, confidence_scores, _ = self.model.run(None, prepared_input) |
| regions = self.postprocess(bboxes, labels, confidence_scores, input_w, input_h) |
| return regions |
|
|
| def postprocess( |
| self, |
| bboxes: np.ndarray, |
| labels: np.ndarray, |
| confidence_scores: np.ndarray, |
| input_w: float, |
| input_h: float, |
| ): |
| """Process output into Unstructured class. Bounding box coordinates are converted to |
| original image resolution.""" |
| regions = [] |
| width_conversion = input_w / self.REQUIRED_WIDTH |
| height_conversion = input_h / self.REQUIRED_HEIGHT |
| for (x1, y1, x2, y2), label, conf in zip(bboxes, labels, confidence_scores): |
| if conf < 0.8: |
| continue |
| region = { |
| "x1": x1 * width_conversion, |
| "y1": y1 * height_conversion, |
| "x2": x2 * width_conversion, |
| "y2": y2 * height_conversion, |
| "confidence": conf, |
| "label": label, |
| } |
|
|
| regions.append(region) |
|
|
| return regions |
|
|
|
|
| @olo.register() |
| def predict(image=olo.File()): |
| model = ONNX_Detectron("./model.onnx") |
| return model.predict(image) |
|
|
|
|
| if __name__ == "__main__": |
| olo.run("detectron") |
|
|
