Spaces:
Sleeping
Sleeping
File size: 2,176 Bytes
59c748e f1495be 59c748e f1495be 59c748e f1495be 59c748e f1495be 59c748e f1495be 59c748e f1495be 59c748e f1495be 59c748e f1495be 59c748e f1495be 59c748e 14c15c7 59c748e f1495be 59c748e 94c7394 |
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 69 70 |
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
|