Spaces:
Runtime error
Runtime error
File size: 1,883 Bytes
6501992 61cc421 6501992 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import cv2
import numpy as np
from ultralytics import YOLO
from PIL import Image
# Load YOLO model
def load_model(model_path="best (10).pt"):
"""Load the YOLO model"""
model = YOLO(model_path)
return model
# Process image
def process_image(image, model):
"""Run YOLO prediction on the given image"""
if isinstance(image, Image.Image):
image_array = np.array(image)
else:
image_array = image
results = model.predict(image_array)
return results[0] # return first result
# Draw all predictions
def draw_predictions(image, results):
"""Draw bounding boxes and labels"""
if isinstance(image, Image.Image):
image_array = np.array(image)
else:
image_array = image
plotted_image = results.plot() # YOLO's built-in plotting
return Image.fromarray(plotted_image)
# Group predictions by class/condition
def group_predictions_by_condition(results):
"""Organize detections by class"""
condition_groups = {}
if len(results.boxes) > 0:
for box in results.boxes:
class_id = int(box.cls[0])
class_name = results.names[class_id]
confidence = float(box.conf[0])
if class_name not in condition_groups:
condition_groups[class_name] = []
condition_groups[class_name].append({
'box': box,
'confidence': confidence
})
return condition_groups
# Example usage (standalone)
if __name__ == "__main__":
model = load_model("best.pt")
img_path = "your_test_image.jpg"
image = Image.open(img_path)
results = process_image(image, model)
grouped = group_predictions_by_condition(results)
print("Grouped Predictions:", grouped)
# Save output visualization
output_img = draw_predictions(image, results)
output_img.save("predicted_output.jpg")
|