def build_detection_engine(model_file: str, labels_file: str):
labels = load_labels(labels_file)
interpreter = tflite.Interpreter(model_path=model_file)
interpreter.allocate_tensors()
_, input_height, input_width, _ = interpreter.get_input_details(
)[0]['shape']
def inference(flow_id: str, frame: object) -> PyDetectionBox:
"""
Returns a list of detection results, each a dictionary of object info
"""
image = frame_data_2_image(frame, input_width, input_height)
set_input_tensor(interpreter, image)
interpreter.invoke()
# Get all output details
boxes = get_output_tensor(interpreter, 0)
classes = get_output_tensor(interpreter, 1)
scores = get_output_tensor(interpreter, 2)
count = int(get_output_tensor(interpreter, 3))
result = PyDetectionBox(frame_id=frame.frame_id, engine_id='tflite')
for i in range(count):
result.add_box(category_id=classes[i],
category_label=labels.get(classes[i], ''),
x1=boxes[i][1],
y1=boxes[i][0],
x2=boxes[i][3],
y2=boxes[i][2],
probability=float(scores[i]))
return flow_id, result
return inference
def build_classification_engine(model_file: str, labels_file: str, top_k=1):
labels = load_labels(labels_file)
interpreter = tflite.Interpreter(model_path=model_file)
interpreter.allocate_tensors()
_, input_height, input_width, _ = interpreter.get_input_details(
)[0]['shape']
def inference(flow_id: str, frame: object) -> PyClassification:
image = frame_data_2_image(frame, input_width, input_height)
set_input_tensor(interpreter, image)
interpreter.invoke()
output_details = interpreter.get_output_details()[0]
output = np.squeeze(interpreter.get_tensor(output_details['index']))
# if the model is quantized (uint8 data), then dequantize the results
if output_details['dtype'] == np.uint8:
scale, zero_point = output_details['quantization']
output = scale * (output - zero_point)
ordered = np.argpartition(-output, top_k)
result = PyClassification(frame_id=frame.frame_id, engine_id='tflite')
for label_id in ordered[:top_k]:
result.add_classification(category_id=label_id,
category_label=labels[label_id],
probability=output[label_id])
return flow_id, result
return inference
def build_detection_engine(model: Callable, labels_file: str):
labels = load_labels(labels_file)
engine = model(pretrained=True)
engine.eval()
transform = get_transform()
def inference(flow_id: str, frame: object) -> PyDetectionBox:
"""
Returns a list of detection results, each a dictionary of object info
"""
image = frame_data_2_image(frame, frame.width, frame.height)
image_t = transform(image)
image_t = torch.unsqueeze(image_t, 0)
output = engine(image_t)
# Get all output details
boxes = output[0]['boxes']
classes = output[0]['labels']
scores = output[0]['scores']
result = PyDetectionBox(frame_id=frame.frame_id, engine_id='tflite')
for i, box in enumerate(boxes):
result.add_box(category_id=classes[i],
category_label=labels.get(classes[i], ''),
x1=box[0],
y1=box[1],
x2=box[2],
y2=box[3],
probability=float(scores[i]))
return flow_id, result
return inference
def build_classification_engine(model_file:str, labels_file:str, top_k:int = 5):
labels = load_labels(labels_file)
session = onnxruntime.InferenceSession(model_file, None)
# get the name of the first input of the model
input_name = session.get_inputs()[0].name
_, input_height, input_width, _ = session.get_inputs()[0].shape
input_type = session.get_inputs()[0].shape
log.info(f'Model input name: {input_name}')
log.info(f'Model input shape: {input_height} x {input_width}')
log.info(f'Model input type: {input_type}')
def score(flow_id:str, frame:object) -> Tuple[str,PyClassification]:
image = frame_data_2_image(frame, input_width, input_height)
image_data = np.array(image).transpose(2, 0, 1)
image_data = preprocess(image_data)
raw_result = session.run([], {input_name:image_data})
res = postprocess(raw_result)
res_idx = np.squeeze(np.argsort(res))[-top_k:]
result = PyClassification(engine_id='onnx_runtime', frame_id=frame.frame_id)
for idx in res_idx:
result.add_classification(category_id=idx, category_label= labels.get(idx, ''), probability=float(res[idx]))
return flow_id, result
return score
def build_classification_engine(model: Callable, labels_file: str, top_k=1):
labels = load_labels(labels_file)
engine = model(pretrained=True)
engine.eval()
transform = get_transform()
def inference(flow_id: str, frame: object) -> PyClassification:
image = frame_data_2_image(frame, frame.width, frame.height)
image_t = transform(image)
image_t = torch.unsqueeze(image_t, 0)
output = engine(image_t)
percentage = torch.nn.functional.softmax(output, dim=1)[0]
_, indices = torch.sort(output, descending=True)
result = PyClassification(frame_id=frame.frame_id,
engine_id='torchvision')
for idx in indices[0][:top_k]:
result.add_classification(category_id=idx,
category_label=labels[idx],
probability=percentage[idx].item())
return flow_id, result
return inference