feat: initial implementation

app.py

@@ -1,94 +1,22 @@
 import gradio as gr
-import torch
+from ultralytics import YOLO
 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 your custom model
+model = YOLO('ethandavey/yoloV5-coursework-model')  # Update with your model's path
 
-# 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()
+    # Run inference
+    results = model(image)
 
-    # Padded resize
-    img = letterbox(img, 640, stride=32, auto=False)[0]
+    # Annotate image
+    annotated_image = results[0].plot()
 
-    # Convert
-    img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
-    img = np.ascontiguousarray(img)
+    return Image.fromarray(annotated_image)
 
-    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)
 
+# Create Gradio Interface
 iface = gr.Interface(
     fn=detect,
     inputs=gr.Image(type="pil"),

requirements.txt

@@ -8,4 +8,4 @@ PyYAML>=5.3.1
 seaborn
 tqdm
 pandas
-git+https://github.com/ultralytics/yolov5.git
+ultralytics