Spaces:
Running
Running
File size: 1,971 Bytes
132787d aae7036 132787d ae1b016 4160d5b 132787d ae1b016 132787d ae1b016 4160d5b ae1b016 aae7036 132787d 4160d5b 132787d 4160d5b ae1b016 4160d5b 132787d 4160d5b 132787d ae1b016 132787d ae1b016 132787d 4160d5b 132787d 4160d5b 132787d |
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 |
from fastapi import FastAPI, UploadFile
from ultralytics import YOLO
from PIL import Image
import os
from huggingface_hub import hf_hub_download
def init_model(model_id: str):
# Define models
MODEL_OPTIONS = {
"YOLOv11-Nano": "medieval-yolov11n.pt",
"YOLOv11-Small": "medieval-yolov11s.pt",
"YOLOv11-Medium": "medieval-yolov11m.pt",
"YOLOv11-Large": "medieval-yolov11l.pt",
"YOLOv11-XLarge": "medieval-yolov11x.pt"
}
if model_id in MODEL_OPTIONS:
print(MODEL_OPTIONS[model_id])
path = hf_hub_download(
repo_id="biglam/medieval-manuscript-yolov11",
filename=MODEL_OPTIONS[model_id],
)
print(path)
# Initialize and return model
model = YOLO(path)
print("Model initialized")
return model
else:
raise ValueError(f"Model {model_id} not found")
app = FastAPI()
@app.get("/")
async def root():
return {"status": "ok"}
@app.post("/predict")
async def predict(image: UploadFile,
model_id: str = "YOLOv11-XLarge",
conf: float = 0.25,
iou: float = 0.7
):
print(model_id, conf, iou)
# Initialize model for each request
model = init_model(model_id)
# Open image from uploaded file
image = Image.open(image.file)
print("Image opened")
# Run inference with the PIL Image
results = model.predict(source=image, conf=conf, iou=iou)
print("Inference done")
# Extract detection results
result = results[0]
detections = []
for box in result.boxes:
detection = {
"class": result.names[int(box.cls[0])],
"confidence": float(box.conf[0]),
"bbox": box.xyxy[0].tolist()
}
detections.append(detection)
return {"detections": detections}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) |