Spaces:
Sleeping
Sleeping
File size: 2,929 Bytes
79fb9b5 885e3d1 79fb9b5 885e3d1 4e2d441 79fb9b5 885e3d1 bc14a59 79fb9b5 4e2d441 79fb9b5 4e2d441 79fb9b5 4e2d441 79fb9b5 4e2d441 79fb9b5 4e2d441 79fb9b5 1780543 4e2d441 79fb9b5 4e2d441 79fb9b5 4e2d441 79fb9b5 4e2d441 79fb9b5 4e2d441 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse, Response, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
import torch
import cv2
import numpy as np
app = FastAPI()
# Разрешаем вызовы из фронта того же Space
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
# Загружаем модель
model = YOLO("best.pt")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
@app.get("/")
def root():
return FileResponse("index.html")
def read_image_to_bgr(file_bytes: bytes) -> np.ndarray:
# Декод JPEG/PNG в BGR
img_array = np.frombuffer(file_bytes, dtype=np.uint8)
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) # BGR
return img
def annotate_bgr(results) -> np.ndarray:
# results[0].plot() возвращает BGR с нарисованными боксами
return results[0].plot()
def results_to_json(results):
# Конвертация результатов в чистые боксы/классы/скор
r = results[0]
boxes = r.boxes
out = []
if boxes is not None and len(boxes) > 0:
xyxy = boxes.xyxy.cpu().numpy() # (N,4)
conf = boxes.conf.cpu().numpy() # (N,)
cls = boxes.cls.cpu().numpy().astype(int) # (N,)
names = r.names
for i in range(len(xyxy)):
x1, y1, x2, y2 = xyxy[i].tolist()
out.append({
"bbox": [x1, y1, x2, y2],
"conf": float(conf[i]),
"class_id": int(cls[i]),
"class_name": names[int(cls[i])] if names else str(cls[i])
})
return {"detections": out}
@app.post("/predict")
async def predict(
file: UploadFile = File(...),
conf: float = Form(0.25),
iou: float = Form(0.45),
return_image: int = Form(1) # 1 = вернуть аннотированное изображение, 0 = вернуть JSON боксов
):
data = await file.read()
bgr = read_image_to_bgr(data)
if bgr is None:
return JSONResponse({"error": "Invalid image"}, status_code=400)
# Инференс (без трекинга — кадры независимы; для трекинга можно persist и tracker)
results = model.predict(
source=bgr,
conf=conf,
iou=iou,
imgsz=640,
verbose=False
)
if return_image == 1:
annotated = annotate_bgr(results) # BGR
# Кодируем в JPEG для отправки
ok, buf = cv2.imencode(".jpg", annotated)
if not ok:
return JSONResponse({"error": "Encode failed"}, status_code=500)
return Response(content=buf.tobytes(), media_type="image/jpeg")
else:
return JSONResponse(results_to_json(results))
|