|
|
from huggingface_hub import hf_hub_download |
|
|
from typing import Dict, List, Any |
|
|
from ultralytics import YOLO |
|
|
import json |
|
|
from urllib.request import urlopen |
|
|
|
|
|
class EndpointHandler(): |
|
|
def __init__(self, path=""): |
|
|
hf_hub_download(repo_id="Drazcat-AI/flejes", filename="yolov8_flejes/runs/detect/train/weights/best.pt") |
|
|
self.model = YOLO(hf_hub_download(repo_id="Drazcat-AI/flejes", filename="yolov8_flejes/runs/detect/train/weights/best.pt", local_files_only=True)) |
|
|
|
|
|
def predict_objects(self, image_path): |
|
|
|
|
|
results = self.model(image_path, imgsz=800) |
|
|
predictions = [] |
|
|
for box in results[0].boxes: |
|
|
class_id = results[0].names[box.cls[0].item()] |
|
|
cords = box.xywh[0].tolist() |
|
|
cords = [round(x) for x in cords] |
|
|
conf = round(box.conf[0].item(), 2) |
|
|
prediction = { |
|
|
"x": cords[0], |
|
|
"y": cords[1], |
|
|
"width": cords[2], |
|
|
"height": cords[3], |
|
|
"confidence": conf, |
|
|
"class": class_id |
|
|
} |
|
|
predictions.append(prediction) |
|
|
predictions_array = {"predictions": predictions} |
|
|
|
|
|
return predictions_array |
|
|
|
|
|
def __call__(self, event): |
|
|
print(event) |
|
|
if "inputs" not in event: |
|
|
return { |
|
|
"statusCode": 400, |
|
|
"body": json.dumps("Error: Please provide an 'inputs' parameter."), |
|
|
} |
|
|
|
|
|
image_path = event["inputs"] |
|
|
|
|
|
try: |
|
|
image = urlopen(image_path).read() |
|
|
predictions = self.predict_objects(image) |
|
|
return { |
|
|
"statusCode": 200, |
|
|
"body": json.dumps(predictions), |
|
|
} |
|
|
except Exception as e: |
|
|
return { |
|
|
"statusCode": 500, |
|
|
"body": json.dumps(f"Error: {str(e)}"), |
|
|
} |
