App

app.py

@@ -0,0 +1,53 @@
+import gradio as gr
+from ultralytics import YOLO
+import cv2
+
+model = YOLO('head_yolov8s.pt')
+
+def head_detect(path):
+    img = cv2.imread(path)
+    output = model(source=img)
+    res = output[0].cpu().numpy()
+    
+    # Extract bbox, cls id and conf
+    bboxes = res.boxes.xyxy
+    class_ids = res.boxes.cls
+    conf_scores = res.boxes.conf
+
+    for i in range(len(bboxes)):
+        xmin, ymin, xmax, ymax = int(bboxes[i][0]), int(bboxes[i][1]), int(bboxes[i][2]), int(bboxes[i][3])
+        conf = conf_scores[i]
+        cls_id = int(class_ids[i])
+        label = model.names[cls_id]  # Get the label name
+
+        # Draw rectangle for bounding box
+        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color=(0, 0, 255), thickness=2, lineType=cv2.LINE_AA)
+
+        # Prepare label text with confidence score
+        label_text = f'{label} {conf:.2f}'
+
+        # Put text (label) on the image
+        cv2.putText(img, label_text, (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, lineType=cv2.LINE_AA)
+
+    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+input_image = [
+    gr.components.Image(type='filepath', label='Input Image'),
+]
+
+output_image = [
+    gr.components.Image(type='numpy', label='Prediction'),
+]
+
+interface = gr.Interface(
+    fn=head_detect,
+    inputs=input_image,
+    outputs=output_image,
+    title='Top Head Detection',
+    cache_examples=False,
+)
+
+gr.TabbedInterface(
+    [interface],
+    tab_names=['Image Inference']
+).queue().launch()