|
|
|
|
|
import os |
|
import cv2 |
|
import numpy as np |
|
import onnxruntime |
|
import pandas as pd |
|
from exif import Image |
|
from pathlib import Path |
|
from utils import load_img, nms, xywh2xyxy |
|
|
|
class ONNXInference(): |
|
def __init__(self, model_path, files, save_image, save_path='./'): |
|
self.model_path = model_path |
|
self.files = files |
|
self.conf_thres = 0.2 |
|
self.iou_thres = 0.7 |
|
self.save_image = save_image |
|
self.save_path = save_path |
|
|
|
def run(self): |
|
opt_session = onnxruntime.SessionOptions() |
|
opt_session.enable_mem_pattern = True |
|
opt_session.enable_cpu_mem_arena = True |
|
opt_session.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL |
|
|
|
EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider'] |
|
ort_session = onnxruntime.InferenceSession(self.model_path, sess_options=opt_session, providers=EP_list) |
|
|
|
model_inputs = ort_session.get_inputs() |
|
input_names = [model_inputs[i].name for i in range(len(model_inputs))] |
|
input_shape = model_inputs[0].shape |
|
print(input_shape) |
|
model_output = ort_session.get_outputs() |
|
output_names = [model_output[i].name for i in range(len(model_output))] |
|
|
|
IMG_ID = [] |
|
PRED_LAB = [] |
|
GEO_TAG_URL = [] |
|
PRED_CT = [] |
|
file_paths = [file.name for file in self.files] |
|
|
|
for i in file_paths: |
|
|
|
image, image_height, image_width, input_height, input_width, input_tensor = load_img( |
|
Path(i).as_posix(), |
|
input_shape |
|
) |
|
|
|
|
|
outputs = ort_session.run(output_names, {input_names[0]: input_tensor})[0] |
|
|
|
predictions = np.squeeze(outputs).T |
|
CONF_THRESHOLD = self.conf_thres |
|
scores = np.max(predictions[:, 4:], axis=1) |
|
predictions = predictions[scores > CONF_THRESHOLD, :] |
|
scores = scores[scores > CONF_THRESHOLD] |
|
class_ids = np.argmax(predictions[:, 4:], axis=1) |
|
boxes = predictions[:, :4] |
|
|
|
|
|
input_shape = np.array([input_width, input_height, input_width, input_height]) |
|
boxes = np.divide(boxes, input_shape, dtype=np.float32) |
|
boxes *= np.array([image_width, image_height, image_width, image_height]) |
|
boxes = boxes.astype(np.int32) |
|
|
|
|
|
|
|
IOU_THRESHOLD = self.iou_thres |
|
indices = nms(xywh2xyxy(boxes), scores, IOU_THRESHOLD) |
|
|
|
if self.save_image: |
|
image_draw = image.copy() |
|
for (bbox, score, label) in zip(xywh2xyxy(boxes[indices]), scores[indices], class_ids[indices]): |
|
bbox = bbox.round().astype(np.int32).tolist() |
|
cls_id = int(label) |
|
CLASSES = ["plastic"] |
|
cls = CLASSES[cls_id] |
|
color = (0,255,0) |
|
cv2.rectangle(image_draw, tuple(bbox[:2]), tuple(bbox[2:]), color, 2) |
|
cv2.putText(image_draw, |
|
f'{cls}:{int(score*100)}', (bbox[0], bbox[1] - 2), |
|
cv2.FONT_HERSHEY_SIMPLEX, |
|
0.60, [225, 255, 255], |
|
thickness=1) |
|
cv2.imwrite(f"result_[0].jpg", image_draw) |
|
|
|
|
|
boxes = boxes[indices] |
|
scores = scores[indices] |
|
pred_ct = len(scores) |
|
|
|
|
|
with open(i, "rb") as image_geo: |
|
my_image = Image(image_geo) |
|
dd_lat = my_image.gps_latitude[0] + (my_image.gps_latitude[1]/60) + (my_image.gps_latitude[2]/3600) |
|
dd_long = my_image.gps_longitude[0] + (my_image.gps_longitude[1]/60) + (my_image.gps_longitude[2]/3600) |
|
url = f"https://www.google.com/maps?q={dd_lat:.7f}%2C{dd_long:.7f}" |
|
GEO_TAG_URL.append(url) |
|
|
|
if pred_ct is not None: |
|
PRED_LAB.append("Yes") |
|
IMG_ID.append(str(i)) |
|
PRED_CT.append(pred_ct) |
|
|
|
|
|
result = { |
|
"IMG_ID": IMG_ID, |
|
"PRED_LAB": PRED_LAB, |
|
"PRED_CT": PRED_CT, |
|
"GEO_TAG_URL": GEO_TAG_URL |
|
} |
|
|
|
return result |