Update app.py

app.py

@@ -1,3 +1,103 @@
+# 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))
+
+
+
+
+
+
+
+
+
+
+
+
+
 from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.responses import JSONResponse, Response, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
@@ -8,7 +108,7 @@ import numpy as np
 
 app = FastAPI()
 
-# Разрешаем вызовы из фронта того же Space
+# Разрешаем фронту обращаться к API
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -26,24 +126,21 @@ 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
+    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
     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,)
+        xyxy = boxes.xyxy.cpu().numpy()
+        conf = boxes.conf.cpu().numpy()
+        cls = boxes.cls.cpu().numpy().astype(int)
         names = r.names
         for i in range(len(xyxy)):
             x1, y1, x2, y2 = xyxy[i].tolist()
@@ -60,26 +157,24 @@ async def predict(
     file: UploadFile = File(...),
     conf: float = Form(0.25),
     iou: float = Form(0.45),
-    return_image: int = Form(1)  # 1 = вернуть аннотированное изображение, 0 = вернуть JSON боксов
+    return_image: int = Form(1)
 ):
     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,
+        imgsz=480,  # меньше размер для скорости
         verbose=False
     )
 
     if return_image == 1:
-        annotated = annotate_bgr(results)  # BGR
-        # Кодируем в JPEG для отправки
-        ok, buf = cv2.imencode(".jpg", annotated)
+        annotated = annotate_bgr(results)
+        ok, buf = cv2.imencode(".jpg", annotated, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
         if not ok:
             return JSONResponse({"error": "Encode failed"}, status_code=500)
         return Response(content=buf.tobytes(), media_type="image/jpeg")