banao-tech commited on
Commit
cc59622
·
verified ·
1 Parent(s): 4b6cfea

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +4 -4
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) -> bool:
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 = Field(0.05, ge=0, le=1),
167
- iou_threshold: float = Field(0.1, ge=0, le=1)
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.