Spaces:
Sleeping
Sleeping
File size: 2,182 Bytes
f1495be |
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 66 67 68 |
# 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
|