Update backend/main.py

backend/main.py

@@ -21,6 +21,7 @@ import ffmpeg
 import torch
 import torchvision.transforms as T
 from ultralytics import YOLO
+import ultralytics.nn.tasks
 import mediapipe as mp
 from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks, Form, Request
 from fastapi.responses import FileResponse, StreamingResponse, JSONResponse
@@ -39,7 +40,8 @@ torch.serialization.add_safe_globals([
     torch.nn.modules.activation.SiLU,
     torch.nn.modules.container.ModuleList,
     torch.nn.modules.upsampling.Upsample,
-    torch.nn.modules.pooling.MaxPool2d
+    torch.nn.modules.pooling.MaxPool2d,
+    ultralytics.nn.tasks.DetectionModel
 ])