File size: 2,050 Bytes
a4c6590 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | 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")
|