Upload 3 files

app.py

@@ -0,0 +1,35 @@
+import gradio as gr
+from predict import ONNXInference
+
+PRED = []
+
+def detect(files):
+    model = ONNXInference(
+        model_path="/home/neo/Downloads/torchFlow/models/torchFlow-ckpt.onnx",
+        files=files,
+        save_image=False,
+        save_path="/home/neo/Downloads/torchFlow/"
+    )
+    res = model.run()
+    img_id = res["IMG_ID"]
+    pred_lab = res["PRED_LAB"],
+    pred_ct = res["PRED_CT"],
+    geo_tag_url =  res["GEO_TAG_URL"]
+    PRED.append(pred_ct)
+    return f"Predicted"
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        output=gr.Image()
+    with gr.Row():
+        btn = gr.UploadButton(
+            label="Upload Image",
+            file_types = ['.jpg','.jpeg'],
+            file_count = "multiple")
+        btn.upload(fn=detect, inputs=btn)
+        with gr.Column(scale=1, min_width=600):
+            gr.Markdown(f"Output here")
+            if PRED is not None:
+                gr.Markdown(f"Predicted: {PRED}")
+
+demo.launch()

predict.py

@@ -0,0 +1,114 @@
+# Inference script for YOLOv8 ONNX model
+
+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 # confidence threshold for onnx model
+        self.iou_thres = 0.7 # intersection-over-union threshold for onnx model
+        self.save_image = save_image
+        self.save_path = save_path
+
+    def run(self):
+        opt_session = onnxruntime.SessionOptions()
+        opt_session.enable_mem_pattern = True # True: memory efficient
+        opt_session.enable_cpu_mem_arena = True # True: memory efficient
+        opt_session.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL # ALL: for optimization
+
+        EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider'] # providers list
+        ort_session = onnxruntime.InferenceSession(self.model_path, sess_options=opt_session, providers=EP_list)
+
+        model_inputs = ort_session.get_inputs() # List of input nodes for loaded ONNX model
+        input_names = [model_inputs[i].name for i in range(len(model_inputs))] # names of the input nodes
+        input_shape = model_inputs[0].shape # shape of input
+        print(input_shape)
+        model_output = ort_session.get_outputs() # list of output nodes for loaded ONNX model
+        output_names = [model_output[i].name for i in range(len(model_output))] # list of output names
+
+        IMG_ID = []
+        PRED_LAB = []
+        GEO_TAG_URL = []
+        PRED_CT = []
+        file_paths = [file.name for file in self.files]
+
+        for i in file_paths:
+            # Loading images
+            image, image_height, image_width, input_height, input_width, input_tensor = load_img(
+                                                                            Path(i).as_posix(),
+                                                                            input_shape
+                                                                        )
+
+            # Run
+            outputs = ort_session.run(output_names, {input_names[0]: input_tensor})[0] # ONNX output as numpy array
+
+            predictions = np.squeeze(outputs).T
+            CONF_THRESHOLD = self.conf_thres
+            scores = np.max(predictions[:, 4:], axis=1)
+            predictions = predictions[scores > CONF_THRESHOLD, :] # Filter out object confidence scores below threshold
+            scores = scores[scores > CONF_THRESHOLD]
+            class_ids = np.argmax(predictions[:, 4:], axis=1)
+            boxes = predictions[:, :4] # (x,y,w,h)
+
+            #rescale box
+            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)
+
+
+            # Apply NMS to suppress weak, overlapping bounding boxes
+            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)
+
+            # preds
+            boxes = boxes[indices]
+            scores = scores[indices]
+            pred_ct = len(scores)
+
+            # Getting geo-coordinates
+            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

utils.py

@@ -0,0 +1,72 @@
+# Utilities
+
+import cv2
+import numpy as np
+from PIL import Image
+
+def load_img(img_path, input_shape):
+    # Loading image
+    image = cv2.imread(img_path)
+    image_height, image_width = image.shape[:2]
+    Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+    input_height, input_width = input_shape[2:]
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    resized = cv2.resize(image_rgb, (input_width, input_height)) # image resized as req by onnx model
+
+    input_image = resized / 255.0 # scaling image
+    input_image = input_image.transpose(2,0,1) # dim rearranged as req by onnx (batch_size, channel, height, width)
+    input_tensor = input_image[np.newaxis, :, :, :].astype(np.float32)
+
+    return image, image_height, image_width, input_height, input_width, input_tensor
+
+def nms(boxes, scores, iou_threshold):
+    # Sort by score
+    sorted_indices = np.argsort(scores)[::-1]
+
+    keep_boxes = []
+    while sorted_indices.size > 0:
+        # Pick the last box
+        box_id = sorted_indices[0]
+        keep_boxes.append(box_id)
+
+        # Compute IoU of the picked box with the rest
+        ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
+
+        # Remove boxes with IoU over the threshold
+        keep_indices = np.where(ious < iou_threshold)[0]
+
+        # print(keep_indices.shape, sorted_indices.shape)
+        sorted_indices = sorted_indices[keep_indices + 1]
+
+    return keep_boxes
+
+def compute_iou(box, boxes):
+    # Compute xmin, ymin, xmax, ymax for both boxes
+    xmin = np.maximum(box[0], boxes[:, 0])
+    ymin = np.maximum(box[1], boxes[:, 1])
+    xmax = np.minimum(box[2], boxes[:, 2])
+    ymax = np.minimum(box[3], boxes[:, 3])
+
+    # Compute intersection area
+    intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
+
+    # Compute union area
+    box_area = (box[2] - box[0]) * (box[3] - box[1])
+    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+    union_area = box_area + boxes_area - intersection_area
+
+    # Compute IoU
+    iou = intersection_area / union_area
+
+    return iou
+
+
+def xywh2xyxy(x):
+    # Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2)
+    y = np.copy(x)
+    y[..., 0] = x[..., 0] - x[..., 2] / 2
+    y[..., 1] = x[..., 1] - x[..., 3] / 2
+    y[..., 2] = x[..., 0] + x[..., 2] / 2
+    y[..., 3] = x[..., 1] + x[..., 3] / 2
+    return y