Update app.py

app.py

@@ -1,12 +1,21 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse, FileResponse
-import cv2
+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"
@@ -16,23 +25,63 @@ model.to(device)
 def root():
     return FileResponse("index.html")
 
-@app.get("/video_feed")
-def video_feed(request: Request):
-    conf = float(request.query_params.get("conf", 0.25))
-    iou = float(request.query_params.get("iou", 0.45))
-    cam_id = int(request.query_params.get("camera", 0))
-
-    cap = cv2.VideoCapture(cam_id)
-
-    def gen_frames():
-        while True:
-            success, frame = cap.read()
-            if not success:
-                break
-            results = model.track(frame, conf=conf, iou=iou, persist=True, tracker="botsort.yaml")
-            annotated = results[0].plot()
-            _, buffer = cv2.imencode(".jpg", annotated)
-            yield (b"--frame\r\n"
-                   b"Content-Type: image/jpeg\r\n\r\n" + buffer.tobytes() + b"\r\n")
-
-    return StreamingResponse(gen_frames(), media_type="multipart/x-mixed-replace; boundary=frame")
+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))