feat: initial implementation

app.py

@@ -0,0 +1,100 @@
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+import cv2
+
+# Import YOLOv5
+from yolov5.models.common import DetectMultiBackend
+from yolov5.utils.torch_utils import select_device
+from yolov5.utils.general import non_max_suppression, scale_coords
+
+# Load model
+device = select_device('cpu')  # use 'cuda:0' for GPU
+model = DetectMultiBackend('ethandavey/yoloV5-coursework-model', device=device)
+model.warmup(imgsz=(1, 3, 640, 640))  # warmup
+
+def detect(image):
+    img = np.array(image)
+    img0 = img.copy()
+
+    # Padded resize
+    img = letterbox(img, 640, stride=32, auto=False)[0]
+
+    # Convert
+    img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
+    img = np.ascontiguousarray(img)
+
+    img = torch.from_numpy(img).to(device)
+    img = img.float()
+    img /= 255.0  # 0 - 255 to 0.0 - 1.0
+    if img.ndimension() == 3:
+        img = img.unsqueeze(0)
+
+    # Inference
+    pred = model(img)
+
+    # NMS
+    pred = non_max_suppression(pred, 0.25, 0.45, classes=None, agnostic=False)
+
+    # Process detections
+    for i, det in enumerate(pred):  # detections per image
+        im0 = img0.copy()
+
+        if len(det):
+            # Rescale boxes from img_size to im0 size
+            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
+
+            for *xyxy, conf, cls in reversed(det):
+                label = f'{model.names[int(cls)]} {conf:.2f}'
+                plot_one_box(xyxy, im0, label=label, color=(255, 0, 0), line_thickness=2)
+
+    return Image.fromarray(im0)
+
+def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), stride=32, auto=True):
+    # Resize and pad image while meeting stride-multiple constraints
+    shape = im.shape[:2]  # current shape [height, width]
+    if isinstance(new_shape, int):
+        new_shape = (new_shape, new_shape)
+
+    # Scale ratio (new / old)
+    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+    r = min(r, 1.0)
+
+    # Compute padding
+    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
+    if auto:  # minimum rectangle
+        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
+
+    dw /= 2  # divide padding into 2 sides
+    dh /= 2
+
+    if shape[::-1] != new_unpad:  # resize
+        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
+    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
+    return im, r, (dw, dh)
+
+def plot_one_box(xyxy, im, color=(128, 128, 128), label=None, line_thickness=3):
+    # Add a bounding box to an image
+    tl = line_thickness or round(0.002 * max(im.shape[0:2])) + 1  # line thickness
+    c1, c2 = (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3]))
+    cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
+    if label:
+        tf = max(tl - 1, 1)  # font thickness
+        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
+        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
+        cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA)  # filled
+        cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, (255, 255, 255), thickness=tf, lineType=cv2.LINE_AA)
+
+iface = gr.Interface(
+    fn=detect,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="YOLOv5 Object Detection",
+    description="Upload an image to detect objects using the YOLOv5 model."
+)
+
+iface.launch()

requirements.txt

@@ -0,0 +1,11 @@
+torch>=1.7.0
+torchvision>=0.8.1
+pillow
+numpy
+matplotlib
+opencv-python
+PyYAML>=5.3.1
+seaborn
+tqdm
+pandas
+git+https://github.com/ultralytics/yolov5.git