Initial commit for YOLOv8 object detection

app.py

@@ -1,70 +1,37 @@
+from ultralytics import YOLO
+from PIL import Image
 import gradio as gr
-import openvino.runtime as ov
-import numpy as np
-from PIL import Image, ImageDraw
-import os
 
-# Define local model path
-MODEL_PATH = "./best_int8_openvino_model/best.xml"  # Path to OpenVINO model file
+# Define the path to the YOLOv8 model
+model_path = "./best_yolo8_model/best.pt"  # Adjust to your local model path
 
+# Function to load the YOLOv8 model
 def load_model(model_path):
-    # Load OpenVINO model
-    print(f"Loading model from {model_path}")
-    core = ov.Core()
-    compiled_model = core.compile_model(model_path, "CPU")
-    return compiled_model
-
-# Load the model
-compiled_model = load_model(MODEL_PATH)
-input_layer = compiled_model.input(0)
-output_layer = compiled_model.output(0)
-
-# Define inference function
-def predict(image):
-    # Preprocess the image
-    image = image.resize((640, 640))  # Resize to model input size
-    image_array = np.array(image).astype(np.float32)
-    image_array = np.transpose(image_array, (2, 0, 1))  # Channels-first for OpenVINO
-    image_array = np.expand_dims(image_array, axis=0)   # Add batch dimension
-
-    # Perform inference
-    results = compiled_model([image_array])
-    output_layer_name = list(results.keys())[0]  # Assuming single output layer
-    boxes = []
-    for result in results[output_layer_name]:
-        box = {
-            'bounding_box': result[:4].tolist(),  # Ensure coordinates are converted to Python list
-            'confidence': float(result[4]),      # Convert to Python float
-            'class_id': int(result[5])           # Convert to Python int
-        }
-        boxes.append(box)
-
-    # Draw bounding boxes
-    draw = ImageDraw.Draw(image)
-    for box in boxes:
-        x_min, y_min, x_max, y_max = map(int, box['bounding_box'])
-
-        # Ensure coordinates are within bounds
-        x_min = max(0, min(x_min, image.width - 1))
-        y_min = max(0, min(y_min, image.height - 1))
-        x_max = max(0, min(x_max, image.width - 1))
-        y_max = max(0, min(y_max, image.height - 1))
-
-        # Draw rectangle and class label
-        draw.rectangle([x_min, y_min, x_max, y_max], outline="red", width=3)
-        draw.text((x_min, y_min), f"Class {box['class_id']}: {box['confidence']:.2f}", fill="red")
-
-    return image
-
-
-# Gradio Interface
-interface = gr.Interface(
-    fn=predict,  # Use the updated predict function
-    inputs=gr.Image(type="pil"),  # Accept an image as input
-    outputs=gr.Image(type="pil"),  # Return the annotated image as output
-    title="OpenVINO Object Detection"  # Title of the app
-)
-
-# Launch the app
-if __name__ == "__main__":
-    interface.launch(share=True)
+    print(f"Loading YOLOv8 model from: {model_path}")
+    detection_model = YOLO(model_path)  # Directly load the local model
+    return detection_model
+
+# Function to perform inference
+def predict(pil_img):
+    try:
+        # Perform prediction
+        results = detection_model.predict(source=pil_img, conf=0.5, iou=0.6, save=False, save_txt=False)
+        
+        # Visualize bounding boxes on the image
+        img_with_boxes = results[0].plot()
+        out_pil_img = Image.fromarray(img_with_boxes[..., ::-1])  # Convert BGR to RGB
+        
+        return out_pil_img
+    except Exception as e:
+        return f"Error during prediction: {str(e)}"
+
+# Load the YOLOv8 model
+detection_model = load_model(model_path)
+
+# Define Gradio interface
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="YOLOv8 Object Detection"
+).launch(share=True)

best_yolo8_model/README.md

@@ -0,0 +1,4 @@
+What the model does.
+Dataset used.
+Training details (e.g., epochs, optimizer).
+License (if required)

best_yolo8_model/best.pt

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