Spaces:
Sleeping
Sleeping
Update main.py
Browse files
main.py
CHANGED
@@ -1,4 +1,4 @@
|
|
1 |
-
from fastapi import FastAPI, File, UploadFile, HTTPException
|
2 |
from fastapi.responses import JSONResponse
|
3 |
from pydantic import BaseModel, Field
|
4 |
from typing import List, Dict, Tuple, Optional
|
@@ -25,7 +25,7 @@ class ModelManager:
|
|
25 |
self.model = None
|
26 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
27 |
|
28 |
-
def load_models(self)
|
29 |
"""Initialize all required models"""
|
30 |
try:
|
31 |
# Load YOLO model
|
@@ -163,8 +163,8 @@ async def startup_event():
|
|
163 |
)
|
164 |
async def process_image(
|
165 |
image_file: UploadFile = File(...),
|
166 |
-
box_threshold: float =
|
167 |
-
iou_threshold: float =
|
168 |
):
|
169 |
"""
|
170 |
Process an uploaded image file and return the results.
|
|
|
1 |
+
from fastapi import FastAPI, File, UploadFile, HTTPException, Query
|
2 |
from fastapi.responses import JSONResponse
|
3 |
from pydantic import BaseModel, Field
|
4 |
from typing import List, Dict, Tuple, Optional
|
|
|
25 |
self.model = None
|
26 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
27 |
|
28 |
+
def load_models(self):
|
29 |
"""Initialize all required models"""
|
30 |
try:
|
31 |
# Load YOLO model
|
|
|
163 |
)
|
164 |
async def process_image(
|
165 |
image_file: UploadFile = File(...),
|
166 |
+
box_threshold: float = Query(0.05, ge=0, le=1),
|
167 |
+
iou_threshold: float = Query(0.1, ge=0, le=1)
|
168 |
):
|
169 |
"""
|
170 |
Process an uploaded image file and return the results.
|