File size: 2,096 Bytes
474ad6f 63474ab 75d7149 474ad6f 589450f 474ad6f 589450f dc18c2d 589450f 474ad6f 2fd8bdc 474ad6f f0775bb 39c617a 474ad6f 589450f 2d78362 589450f 2d78362 700fb92 474ad6f |
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 |
from huggingface_hub import hf_hub_download
from typing import Dict, List, Any
from ultralytics import YOLO
import json
from urllib.request import urlopen
from PIL import Image
from io import BytesIO
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=1280)
#results = self.model(image_path, imgsz=800)
#results = self.model(image_path)
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 = box.conf[0].item()
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):
if "inputs" not in event:
return {
"statusCode": 400,
"body": json.dumps("Error: Please provide an 'inputs' parameter."),
}
image_path = event["inputs"]
try:
#url image
#data=urlopen(image_path).read()
#image = Image.open(BytesIO(data))
#bytes image
image=image_path
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)}"),
} |