initial commit

app.py

@@ -0,0 +1,99 @@
+from typing import Any
+import cv2
+import numpy as np
+from yolov7 import YOLOv7
+import gradio as gr
+from PIL import Image
+
+class Inference:
+    def setup_models(self, model_path, labels_path, engine_path):
+        yolo = YOLOv7(
+            model_path,
+            labels_path,
+            engine_path
+        )
+        return yolo
+    
+    def __init__(self, model_path, labels_path, engine_path):
+        self.model = self.setup_models(
+            model_path,
+            labels_path,
+            engine_path
+        )
+
+    def __call__(self, frame: np.ndarray, conf_threshold: float, nms_threshold: float, *args: Any, **kwds: Any) -> Any:
+        boxes, scores, class_ids = self.model(frame, conf_threshold, nms_threshold)
+        return boxes, scores, class_ids
+
+infer1 = Inference(
+    "models/firesmoke.onnx",
+    "models/labels.txt",
+    "firesmoke.trt"
+)
+
+infer2 = Inference(
+    "models/firesmoke-henry.onnx",
+    "models/labels.txt",
+    "firesmoke-henry.trt"
+)
+
+def run(content_img, conf_threshold, nms_threshold):
+    content_img = cv2.cvtColor(np.array(content_img), cv2.COLOR_RGB2BGR)
+
+    boxes1, scores1, class_ids1 = infer1(content_img, conf_threshold, nms_threshold)
+    boxes2, scores2, class_ids2 = infer2(content_img, conf_threshold, nms_threshold)
+    img1 = content_img.copy()
+    img2 = content_img.copy()
+
+    if len(boxes1) > 0:
+        for box, score, class_id in zip(boxes1, scores1, class_ids1):
+            x1 = int(box[0])
+            y1 = int(box[1])
+            x2 = int(box[2])
+            y2 = int(box[3])
+            cv2.rectangle(img1, (x1, y1), (x2, y2), (0, 0, 255), 2)
+            cv2.rectangle(img1, (x1, y1-20), (x1+100, y1), (0, 0, 255), -1)
+            cv2.putText(img1, "{}:{:.2f}".format(class_id, score), (x1, y1), cv2.FONT_HERSHEY_COMPLEX, 0.5, (255, 255, 255), 1)
+
+    if len(boxes2) > 0:
+        for box, score, class_id in zip(boxes2, scores2, class_ids2):
+            x1 = int(box[0])
+            y1 = int(box[1])
+            x2 = int(box[2])
+            y2 = int(box[3])
+            cv2.rectangle(img2, (x1, y1), (x2, y2), (0, 0, 255), 2)
+            cv2.rectangle(img2, (x1, y1-20), (x1+100, y1), (0, 0, 255), -1)
+            cv2.putText(img2, "{}:{:.2f}".format(class_id, score), (x1, y1), cv2.FONT_HERSHEY_COMPLEX, 0.5, (255, 255, 255), 1)
+
+    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
+    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
+    img1 = Image.fromarray(img1)
+    img2 = Image.fromarray(img2)
+
+    return img1, img2
+
+if __name__ == '__main__':
+    style = gr.Interface(
+        fn=run, 
+        inputs=[
+            gr.Image(label='Input Image'),
+            gr.Slider(minimum=0.05, maximum=1, step=0.05, default=0.3, label="Confidence Threshold"),
+            gr.Slider(minimum=0.05, maximum=1, step=0.05, default=0.3, label="NMS Threshold"),
+        ], 
+        outputs=[
+            gr.Image(
+                type="pil",
+                label="Finetuned"
+            ),
+            gr.Image(
+                type="pil",
+                label="Finetuned + New Data"
+            ),
+        ], 
+        examples=[
+            ['examples/fire1.jpg'],
+            ['examples/fire2.jpg'],
+            ['examples/fire3.jpg']
+        ]
+    )
+    style.launch()

models/firesmoke-henry.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6223c3d99acb5515abce9056de3b5151df5bcc321b7927815c325f79e95420c1
+size 146024150

models/firesmoke.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55d1a9543bc6f610c0d9d6ec1ad64ce3df0ff703f29a92325480fa4f21f807f7
+size 146024150

models/labels.txt

@@ -0,0 +1,2 @@
+smoke
+fire

requirements.txt

@@ -0,0 +1,5 @@
+gradio==3.50.2
+numpy==1.26.1
+onnxruntime==1.16.1
+opencv_python==4.8.1.78
+Pillow==10.1.0

yolov7.py

@@ -0,0 +1,198 @@
+import cv2
+import numpy as np
+import onnxruntime
+
+class YOLOv7:
+
+    def __init__(
+            self, 
+            model_path, 
+            labels_path,
+            engine_path, 
+            official_nms=False
+        ):
+        self.official_nms = official_nms
+        
+        self.class_names = []
+        with open(labels_path, 'r') as f:
+            self.class_names = [cname.strip() for cname in f.readlines()]
+        f.close()
+
+        # Create a list of colors for each class where each color is a tuple of 3 integer values
+        rng = np.random.default_rng(3)
+        self.colors = rng.uniform(0, 255, size=(len(self.class_names), 3))
+
+        # Initialize model
+        self.initialize_model(model_path, engine_path)
+
+    def __call__(self, image, confidence_threshold, nms_threshold):
+        return self.detect_objects(image, confidence_threshold, nms_threshold)
+
+    def xywh2xyxy(self, 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
+
+    def initialize_model(self, model_path, engine_path):
+        self.session = onnxruntime.InferenceSession(
+            model_path,
+            providers=[
+                # (
+                #     'TensorrtExecutionProvider',
+                #     {
+                #         'device_id': 0,
+                #         'trt_max_workspace_size': 2147483648,
+                #         'trt_fp16_enable': True,
+                #         'trt_engine_cache_enable': True,
+                #         'trt_engine_cache_path': '{}'.format(engine_path),
+                #     }
+                # ),
+                # (
+                #     'CUDAExecutionProvider', 
+                #     {
+                #         'device_id': 0,
+                #         'arena_extend_strategy': 'kNextPowerOfTwo',
+                #         'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
+                #         'cudnn_conv_algo_search': 'EXHAUSTIVE',
+                #         'do_copy_in_default_stream': True,
+                #     }
+                # )
+                'CPUExecutionProvider'
+            ]
+        )
+        # Get model info
+        self.get_input_details()
+        self.get_output_details()
+
+        self.has_postprocess = 'score' in self.output_names or self.official_nms
+
+
+    def detect_objects(self, image, confidence_threshold, nms_threshold):
+        input_tensor = self.prepare_input(image)
+
+        # Perform inference on the image
+        outputs = self.inference(input_tensor)
+
+        # Process output data
+        self.boxes, self.scores, self.class_ids = self.process_output(outputs, confidence_threshold, nms_threshold)
+
+        return self.boxes, self.scores, self.class_ids
+
+    def prepare_input(self, image):
+        self.img_height, self.img_width = image.shape[:2]
+
+        input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        # Resize input image
+        input_img = cv2.resize(input_img, (self.input_width, self.input_height))
+
+        # Scale input pixel values to 0 to 1
+        input_img = input_img / 255.0
+        input_img = input_img.transpose(2, 0, 1)
+        input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32)
+
+        return input_tensor
+    
+    def rescale_boxes(self, boxes):
+
+        # Rescale boxes to original image dimensions
+        input_shape = np.array([self.input_width, self.input_height, self.input_width, self.input_height])
+        boxes = np.divide(boxes, input_shape, dtype=np.float32)
+        boxes *= np.array([self.img_width, self.img_height, self.img_width, self.img_height])
+        return boxes
+    
+    def process_output(self, output, conf_threshold, nms_threshold):
+        boxes, scores, class_ids = output
+        boxes = boxes[0]
+        scores = scores[0]
+        class_ids = class_ids[0]
+
+        res_boxes = []
+        res_scores = []
+        res_class_ids = []
+
+        for box, score, class_id in zip(boxes, scores, class_ids):
+            if score > conf_threshold:
+                score = score[0]
+                res_boxes.append(box)
+                res_scores.append(score)
+                res_class_ids.append(int(class_id))
+        
+        if len(res_scores) == 0:
+            return [], [], []
+        
+        # Scale boxes to original image dimensions
+        res_boxes = self.rescale_boxes(res_boxes)
+
+        fin_boxes, fin_scores, fin_class_ids = [], [], []
+        final_boxes = cv2.dnn.NMSBoxes(res_boxes, res_scores, conf_threshold, nms_threshold)
+        for max_valueid in final_boxes:
+            fin_boxes.append(res_boxes[max_valueid])
+            fin_scores.append(res_scores[max_valueid])
+            fin_class_ids.append(res_class_ids[max_valueid])
+        
+        # Convert boxes to xyxy format
+        fin_boxes = self.xywh2xyxy(np.array(fin_boxes))
+        
+        # Convert class ids to class names
+        fin_class_ids = [self.class_names[i] for i in fin_class_ids]
+        return fin_boxes, fin_scores, fin_class_ids
+
+    def draw_detections(self, image, draw_scores=True, mask_alpha=0.4):
+        
+        mask_img = image.copy()
+        det_img = image.copy()
+
+        img_height, img_width = image.shape[:2]
+        size = min([img_height, img_width]) * 0.0006
+        text_thickness = int(min([img_height, img_width]) * 0.001)
+
+        # Draw bounding boxes and labels of detections
+        for box, score, class_id in zip(self.boxes, self.scores, self.class_ids):
+            color = self.colors[class_id]
+
+            x1, y1, x2, y2 = box.astype(int)
+
+            # Draw rectangle
+            cv2.rectangle(det_img, (x1, y1), (x2, y2), color, 2)
+
+            # Draw fill rectangle in mask image
+            cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, -1)
+
+            label = self.class_names[class_id]
+            caption = f'{label} {int(score * 100)}%'
+            (tw, th), _ = cv2.getTextSize(text=caption, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                                        fontScale=size, thickness=text_thickness)
+            th = int(th * 1.2)
+
+            cv2.rectangle(det_img, (x1, y1),
+                        (x1 + tw, y1 - th), color, -1)
+            cv2.rectangle(mask_img, (x1, y1),
+                        (x1 + tw, y1 - th), color, -1)
+            cv2.putText(det_img, caption, (x1, y1),
+                        cv2.FONT_HERSHEY_SIMPLEX, size, (255, 255, 255), text_thickness, cv2.LINE_AA)
+
+            cv2.putText(mask_img, caption, (x1, y1),
+                        cv2.FONT_HERSHEY_SIMPLEX, size, (255, 255, 255), text_thickness, cv2.LINE_AA)
+
+        return cv2.addWeighted(mask_img, mask_alpha, det_img, 1 - mask_alpha, 0)
+
+    def get_input_details(self):
+        model_inputs = self.session.get_inputs()
+        self.input_names = [model_inputs[i].name for i in range(len(model_inputs))]
+
+        self.input_shape = model_inputs[0].shape
+        self.input_height = self.input_shape[2]
+        self.input_width = self.input_shape[3]
+
+    def get_output_details(self):
+        model_outputs = self.session.get_outputs()
+        self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]
+
+    def inference(self, input_tensor):
+        outputs = self.session.run(self.output_names, {self.input_names[0]: input_tensor})
+        return outputs