update

app.py

@@ -1,73 +1,179 @@
-import gradio as gr
-#import torch
-import yolov7
+# import gradio as gr
+# #import torch
+# import yolov7
 
 
-#
-# from huggingface_hub import hf_hub_download
-from huggingface_hub import HfApi
+# #
+# # from huggingface_hub import hf_hub_download
+# from huggingface_hub import HfApi
 
 
-# Images
-#torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg')
-#torch.hub.download_url_to_file('https://raw.githubusercontent.com/obss/sahi/main/tests/data/small-vehicles1.jpeg', 'small-vehicles1.jpeg')
+# # Images
+# #torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg')
+# #torch.hub.download_url_to_file('https://raw.githubusercontent.com/obss/sahi/main/tests/data/small-vehicles1.jpeg', 'small-vehicles1.jpeg')
     
-def yolov7_inference(
-    image: gr.inputs.Image = None,
-    model_path: gr.inputs.Dropdown = None,
-    image_size: gr.inputs.Slider = 640,
-    conf_threshold: gr.inputs.Slider = 0.25,
-    iou_threshold: gr.inputs.Slider = 0.45,
-):
-    """
-    YOLOv7 inference function
-    Args:
-        image: Input image
-        model_path: Path to the model
-        image_size: Image size
-        conf_threshold: Confidence threshold
-        iou_threshold: IOU threshold
-    Returns:
-        Rendered image
-    """
-
-    model = yolov7.load(model_path, device="cpu", hf_model=True, trace=False)
-    model.conf = conf_threshold
-    model.iou = iou_threshold
-    results = model([image], size=image_size)
-    return results.render()[0]
+# def yolov7_inference(
+#     image: gr.inputs.Image = None,
+#     model_path: gr.inputs.Dropdown = None,
+#     image_size: gr.inputs.Slider = 640,
+#     conf_threshold: gr.inputs.Slider = 0.25,
+#     iou_threshold: gr.inputs.Slider = 0.45,
+# ):
+#     """
+#     YOLOv7 inference function
+#     Args:
+#         image: Input image
+#         model_path: Path to the model
+#         image_size: Image size
+#         conf_threshold: Confidence threshold
+#         iou_threshold: IOU threshold
+#     Returns:
+#         Rendered image
+#     """
+
+#     model = yolov7.load(model_path, device="cpu", hf_model=True, trace=False)
+#     model.conf = conf_threshold
+#     model.iou = iou_threshold
+#     results = model([image], size=image_size)
+#     return results.render()[0]
         
 
 
-inputs = [
-    gr.inputs.Image(type="pil", label="Input Image"),
-    gr.inputs.Dropdown(
-        choices=[
-            "alshimaa/model_baseline",
-            "alshimaa/model_yolo7",
-            #"kadirnar/yolov7-v0.1",
-        ],
-        default="alshimaa/model_baseline",
-        label="Model",
+# inputs = [
+#     gr.inputs.Image(type="pil", label="Input Image"),
+#     gr.inputs.Dropdown(
+#         choices=[
+#             "alshimaa/model_baseline",
+#             "alshimaa/model_yolo7",
+#             #"kadirnar/yolov7-v0.1",
+#         ],
+#         default="alshimaa/model_baseline",
+#         label="Model",
+#     )
+#     #gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size")
+#     #gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
+#     #gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold")
+# ]
+
+# outputs = gr.outputs.Image(type="filepath", label="Output Image")
+# title = "Smart Environmental Eye (SEE)"
+
+# examples = [['image1.jpg', 'alshimaa/model_yolo7', 640, 0.25, 0.45], ['image2.jpg', 'alshimaa/model_yolo7', 640, 0.25, 0.45], ['image3.jpg', 'alshimaa/model_yolo7', 640, 0.25, 0.45]]
+# demo_app = gr.Interface(
+#     fn=yolov7_inference,
+#     inputs=inputs,
+#     outputs=outputs,
+#     title=title,
+#     examples=examples,
+#     cache_examples=True,
+#     theme='huggingface',
+# )
+
+# demo_app.launch(debug=True, enable_queue=True)
+
+
+
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+
+import cv2
+import gradio as gr
+
+from inferer import Inferer
+
+pipeline = Inferer("alshimaa/model_baseline", device='cuda')
+
+
+def fn_image(image, conf_thres, iou_thres):
+    return pipeline(image, conf_thres, iou_thres)
+
+
+def fn_video(video_file, conf_thres, iou_thres, start_sec, duration):
+    start_timestamp = time.strftime("%H:%M:%S", time.gmtime(start_sec))
+    end_timestamp = time.strftime("%H:%M:%S", time.gmtime(start_sec + duration))
+
+    suffix = Path(video_file).suffix
+
+    clip_temp_file = tempfile.NamedTemporaryFile(suffix=suffix)
+    subprocess.call(
+        f"ffmpeg -y -ss {start_timestamp} -i {video_file} -to {end_timestamp} -c copy {clip_temp_file.name}".split()
     )
-    #gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size")
-    #gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
-    #gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold")
-]
-
-outputs = gr.outputs.Image(type="filepath", label="Output Image")
-title = "Smart Environmental Eye (SEE)"
-
-examples = [['image1.jpg', 'alshimaa/model_yolo7', 640, 0.25, 0.45], ['image2.jpg', 'alshimaa/model_yolo7', 640, 0.25, 0.45], ['image3.jpg', 'alshimaa/model_yolo7', 640, 0.25, 0.45]]
-demo_app = gr.Interface(
-    fn=yolov7_inference,
-    inputs=inputs,
-    outputs=outputs,
-    title=title,
-    examples=examples,
-    cache_examples=True,
-    theme='huggingface',
+
+    # Reader of clip file
+    cap = cv2.VideoCapture(clip_temp_file.name)
+
+    # This is an intermediary temp file where we'll write the video to
+    # Unfortunately, gradio doesn't play too nice with videos rn so we have to do some hackiness
+    # with ffmpeg at the end of the function here.
+    with tempfile.NamedTemporaryFile(suffix=".mp4") as temp_file:
+        out = cv2.VideoWriter(temp_file.name, cv2.VideoWriter_fourcc(*"MP4V"), 30, (1280, 720))
+
+        num_frames = 0
+        max_frames = duration * 30
+        while cap.isOpened():
+            try:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+            except Exception as e:
+                print(e)
+                continue
+            print("FRAME DTYPE", type(frame))
+            out.write(pipeline(frame, conf_thres, iou_thres))
+            num_frames += 1
+            print("Processed {} frames".format(num_frames))
+            if num_frames == max_frames:
+                break
+
+        out.release()
+
+        # Aforementioned hackiness
+        out_file = tempfile.NamedTemporaryFile(suffix="out.mp4", delete=False)
+        subprocess.run(f"ffmpeg -y -loglevel quiet -stats -i {temp_file.name} -c:v libx264 {out_file.name}".split())
+
+    return out_file.name
+
+
+image_interface = gr.Interface(
+    fn=fn_image,
+    inputs=[
+        "image",
+        gr.Slider(0, 1, value=0.5, label="Confidence Threshold"),
+        gr.Slider(0, 1, value=0.5, label="IOU Threshold"),
+    ],
+    outputs=gr.Image(type="file"),
+    examples=[["image1.jpg", 0.5, 0.5], ["image2.jpg", 0.25, 0.45], ["image3.jpg", 0.25, 0.45]],
+    title="Smart Environmental Eye (SEE)",
+    allow_flagging=False,
+    allow_screenshot=False,
 )
 
-demo_app.launch(debug=True, enable_queue=True)
+video_interface = gr.Interface(
+    fn=fn_video,
+    inputs=[
+        gr.Video(type="file"),
+        gr.Slider(0, 1, value=0.25, label="Confidence Threshold"),
+        gr.Slider(0, 1, value=0.45, label="IOU Threshold"),
+        gr.Slider(0, 10, value=0, label="Start Second", step=1),
+        gr.Slider(0, 10 if pipeline.device.type != 'cpu' else 3, value=4, label="Duration", step=1),
+    ],
+    outputs=gr.Video(type="file", format="mp4"),
+    # examples=[
+    #     ["video.mp4", 0.25, 0.45, 0, 2],
+       
+    # ],
+    title="Smart Environmental Eye (SEE)",
+    allow_flagging=False,
+    allow_screenshot=False,
+)
+
+
+
+if __name__ == "__main__":
+    gr.TabbedInterface(
+        [image_interface, video_interface],
+        ["Run on Images", "Run on Videos"],
+    ).launch()
 

inferer.py

@@ -0,0 +1,238 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+import math
+import os.path as osp
+
+import cv2
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image, ImageFont
+
+from yolov6.data.data_augment import letterbox
+from yolov6.layers.common import DetectBackend
+from yolov6.utils.events import LOGGER, load_yaml
+from yolov6.utils.nms import non_max_suppression
+
+
+class Inferer:
+    def __init__(self, model_id, device="cpu", yaml="coco.yaml", img_size=640, half=False):
+        self.__dict__.update(locals())
+
+        # Init model
+        self.img_size = img_size
+        cuda = device != "cpu" and torch.cuda.is_available()
+        self.device = torch.device("cuda:0" if cuda else "cpu")
+        self.model = DetectBackend(hf_hub_download(model_id, "model.pt"), device=self.device)
+        self.stride = self.model.stride
+        self.class_names = load_yaml(yaml)["names"]
+        self.img_size = self.check_img_size(self.img_size, s=self.stride)  # check image size
+
+        # Half precision
+        if half & (self.device.type != "cpu"):
+            self.model.model.half()
+        else:
+            self.model.model.float()
+            half = False
+
+        if self.device.type != "cpu":
+            self.model(
+                torch.zeros(1, 3, *self.img_size).to(self.device).type_as(next(self.model.model.parameters()))
+            )  # warmup
+
+        # Switch model to deploy status
+        self.model_switch(self.model, self.img_size)
+
+    def model_switch(self, model, img_size):
+        """Model switch to deploy status"""
+        from yolov6.layers.common import RepVGGBlock
+
+        for layer in model.modules():
+            if isinstance(layer, RepVGGBlock):
+                layer.switch_to_deploy()
+
+        LOGGER.info("Switch model to deploy modality.")
+
+    def __call__(
+        self,
+        path_or_image,
+        conf_thres=0.25,
+        iou_thres=0.45,
+        classes=None,
+        agnostic_nms=False,
+        max_det=1000,
+        hide_labels=False,
+        hide_conf=False,
+    ):
+        """Model Inference and results visualization"""
+
+        img, img_src = self.precess_image(path_or_image, self.img_size, self.stride, self.half)
+        img = img.to(self.device)
+        if len(img.shape) == 3:
+            img = img[None]
+            # expand for batch dim
+        pred_results = self.model(img)
+        det = non_max_suppression(pred_results, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)[0]
+
+        gn = torch.tensor(img_src.shape)[[1, 0, 1, 0]]  # normalization gain whwh
+        img_ori = img_src
+
+        # check image and font
+        assert (
+            img_ori.data.contiguous
+        ), "Image needs to be contiguous. Please apply to input images with np.ascontiguousarray(im)."
+        self.font_check()
+
+        if len(det):
+            det[:, :4] = self.rescale(img.shape[2:], det[:, :4], img_src.shape).round()
+
+            for *xyxy, conf, cls in reversed(det):
+                class_num = int(cls)  # integer class
+                label = (
+                    None
+                    if hide_labels
+                    else (self.class_names[class_num] if hide_conf else f"{self.class_names[class_num]} {conf:.2f}")
+                )
+
+                self.plot_box_and_label(
+                    img_ori,
+                    max(round(sum(img_ori.shape) / 2 * 0.003), 2),
+                    xyxy,
+                    label,
+                    color=self.generate_colors(class_num, True),
+                )
+
+            img_src = np.asarray(img_ori)
+
+            return img_src
+
+    @staticmethod
+    def precess_image(path_or_image, img_size, stride, half):
+        """Process image before image inference."""
+        if isinstance(path_or_image, str):
+            try:
+                img_src = cv2.imread(path_or_image)
+                assert img_src is not None, f"Invalid image: {path_or_image}"
+            except Exception as e:
+                LOGGER.warning(e)
+        elif isinstance(path_or_image, np.ndarray):
+            img_src = path_or_image
+        elif isinstance(path_or_image, Image.Image):
+            img_src = np.array(path_or_image)
+
+        image = letterbox(img_src, img_size, stride=stride)[0]
+
+        # Convert
+        image = image.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
+        image = torch.from_numpy(np.ascontiguousarray(image))
+        image = image.half() if half else image.float()  # uint8 to fp16/32
+        image /= 255  # 0 - 255 to 0.0 - 1.0
+
+        return image, img_src
+
+    @staticmethod
+    def rescale(ori_shape, boxes, target_shape):
+        """Rescale the output to the original image shape"""
+        ratio = min(ori_shape[0] / target_shape[0], ori_shape[1] / target_shape[1])
+        padding = (ori_shape[1] - target_shape[1] * ratio) / 2, (ori_shape[0] - target_shape[0] * ratio) / 2
+
+        boxes[:, [0, 2]] -= padding[0]
+        boxes[:, [1, 3]] -= padding[1]
+        boxes[:, :4] /= ratio
+
+        boxes[:, 0].clamp_(0, target_shape[1])  # x1
+        boxes[:, 1].clamp_(0, target_shape[0])  # y1
+        boxes[:, 2].clamp_(0, target_shape[1])  # x2
+        boxes[:, 3].clamp_(0, target_shape[0])  # y2
+
+        return boxes
+
+    def check_img_size(self, img_size, s=32, floor=0):
+        """Make sure image size is a multiple of stride s in each dimension, and return a new shape list of image."""
+        if isinstance(img_size, int):  # integer i.e. img_size=640
+            new_size = max(self.make_divisible(img_size, int(s)), floor)
+        elif isinstance(img_size, list):  # list i.e. img_size=[640, 480]
+            new_size = [max(self.make_divisible(x, int(s)), floor) for x in img_size]
+        else:
+            raise Exception(f"Unsupported type of img_size: {type(img_size)}")
+
+        if new_size != img_size:
+            print(f"WARNING: --img-size {img_size} must be multiple of max stride {s}, updating to {new_size}")
+        return new_size if isinstance(img_size, list) else [new_size] * 2
+
+    def make_divisible(self, x, divisor):
+        # Upward revision the value x to make it evenly divisible by the divisor.
+        return math.ceil(x / divisor) * divisor
+
+    @staticmethod
+    def plot_box_and_label(image, lw, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255)):
+        # Add one xyxy box to image with label
+        p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
+        cv2.rectangle(image, p1, p2, color, thickness=lw, lineType=cv2.LINE_AA)
+        if label:
+            tf = max(lw - 1, 1)  # font thickness
+            w, h = cv2.getTextSize(label, 0, fontScale=lw / 3, thickness=tf)[0]  # text width, height
+            outside = p1[1] - h - 3 >= 0  # label fits outside box
+            p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
+            cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA)  # filled
+            cv2.putText(
+                image,
+                label,
+                (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
+                0,
+                lw / 3,
+                txt_color,
+                thickness=tf,
+                lineType=cv2.LINE_AA,
+            )
+
+    @staticmethod
+    def font_check(font="./yolov6/utils/Arial.ttf", size=10):
+        # Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary
+        assert osp.exists(font), f"font path not exists: {font}"
+        try:
+            return ImageFont.truetype(str(font) if font.exists() else font.name, size)
+        except Exception as e:  # download if missing
+            return ImageFont.truetype(str(font), size)
+
+    @staticmethod
+    def box_convert(x):
+        # Convert boxes with shape [n, 4] from [x1, y1, x2, y2] to [x, y, w, h] where x1y1=top-left, x2y2=bottom-right
+        y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+        y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
+        y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
+        y[:, 2] = x[:, 2] - x[:, 0]  # width
+        y[:, 3] = x[:, 3] - x[:, 1]  # height
+        return y
+
+    @staticmethod
+    def generate_colors(i, bgr=False):
+        hex = (
+            "FF3838",
+            "FF9D97",
+            "FF701F",
+            "FFB21D",
+            "CFD231",
+            "48F90A",
+            "92CC17",
+            "3DDB86",
+            "1A9334",
+            "00D4BB",
+            "2C99A8",
+            "00C2FF",
+            "344593",
+            "6473FF",
+            "0018EC",
+            "8438FF",
+            "520085",
+            "CB38FF",
+            "FF95C8",
+            "FF37C7",
+        )
+        palette = []
+        for iter in hex:
+            h = "#" + iter
+            palette.append(tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4)))
+        num = len(palette)
+        color = palette[int(i) % num]
+        return (color[2], color[1], color[0]) if bgr else color