model/pothole_model.py

# from ultralyticsplus import YOLO, render_result
# from PIL import Image
# 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((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
from transformers import AutoImageProcessor, AutoModelForImageClassification

# Load model
processor = AutoImageProcessor.from_pretrained("taroii/pothole-detection-model")
model = AutoModelForImageClassification.from_pretrained("taroii/pothole-detection-model")

# 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