Spaces:
Sleeping
Sleeping
| 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 | |