Create identification_model.py

identification_model.py

@@ -0,0 +1,35 @@
+import torch
+from ultralytics import YOLO
+from PIL import Image
+import numpy as np
+
+def load_yolov8_model():
+    # Load the YOLOv8 model
+    model = YOLO('yolov8n.pt')  # Using the smallest version for speed; adjust as needed
+    return model
+
+def run_object_detection(model, image_path):
+    # Load the image
+    image = Image.open(image_path).convert('RGB')
+    image_np = np.array(image)
+    
+    # Run inference
+    results = model(image_np)
+    
+    # Process results
+    detections = []
+    for result in results:
+        boxes = result.boxes
+        for box in boxes:
+            x1, y1, x2, y2 = box.xyxy[0].tolist()  # Bounding box coordinates
+            conf = box.conf[0].item()  # Confidence score
+            cls = int(box.cls[0].item())  # Class ID
+            label = model.names[cls]  # Class name
+
+            detections.append({
+                'box': [x1, y1, x2, y2],
+                'confidence': conf,
+                'label': label
+            })
+    
+    return detections