from ultralyticsplus import YOLO, render_result
from PIL import Image
from io import BytesIO
import numpy as np

def load_model(image):
    # image_bytes = image.content
    model = YOLO('keremberke/yolov8n-pothole-segmentation')
    model.overrides['conf'] = 0.25
    model.overrides['iou'] = 0.45
    model.overrides['agnostic_nms'] = False
    model.overrides['max_det'] = 1000

    # Load image using PIL
    image = Image.open(BytesIO(image))
    image_array = np.array(image)
    # pil_image = pil_image.convert("RGB")  # Ensure image is in RGB format

    # Convert PIL image to bytes
    # with io.BytesIO() as output:
    #     pil_image.save(output, format='JPEG')
    #     image_bytes = output.getvalue()

    results = model.predict(image_array)
    for result in results:
        boxes = result.boxes.xyxy
        conf = result.boxes.conf
        cls = result.boxes.cls
        obj_info = []
        for i, bbox in enumerate(boxes):
            label = result.names[int(cls[i])]
            obj_info.append({
                "Object": i+1,
                "Label": label,
                "Confidence": conf[i],
                "Bounding Box": bbox
            })
    render = render_result(model=model, image=image, result=results[0])
    if label:
        print(label)
    render.show()
    return label


# from PIL import Image
# from io import BytesIO

# # Load model directly
# from transformers import AutoImageProcessor, AutoModelForObjectDetection

# processor = AutoImageProcessor.from_pretrained("savioratharv/pothole_detection")
# model = AutoModelForObjectDetection.from_pretrained("savioratharv/pothole_detection")

# # Function to predict if an image contains a pothole
# def predict_pothole(image_url):
#     image = Image.open(BytesIO(image_url))
#     inputs = processor(images=image, return_tensors="pt")

#     # Perform inference
#     outputs = model(**inputs)
#     logits = outputs.logits
#     probabilities = logits.softmax(dim=1)
    
#     # Get predicted class (0: No pothole, 1: Pothole)
#     predicted_class = probabilities.argmax().item()
#     confidence = probabilities[0, predicted_class].item()

#     return predicted_class