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")