satvs commited on
Commit
d476ed4
·
1 Parent(s): 81d3fd3

[ref] update to pruned model

Browse files
Files changed (2) hide show
  1. app.py +2 -2
  2. tasks/image.py +2 -1
app.py CHANGED
@@ -16,8 +16,8 @@ app.include_router(image.router)
16
  @app.get("/")
17
  async def root():
18
  return {
19
- "message": "This is an API for a frugal object detection model",
20
  "endpoints": {
21
- "image": "/image - Image classification task (coming soon)",
22
  }
23
  }
 
16
  @app.get("/")
17
  async def root():
18
  return {
19
+ "message": "This is an API for a frugal fire smoke detection model",
20
  "endpoints": {
21
+ "image": "/image - Image classification task",
22
  }
23
  }
tasks/image.py CHANGED
@@ -103,7 +103,8 @@ async def evaluate_image(request: ImageEvaluationRequest):
103
 
104
  # Load model
105
  model_path = Path("tasks", "models")
106
- model_name = "best.pt"
 
107
  model = YOLO(Path(model_path, model_name))
108
  threshold = 0.14
109
 
 
103
 
104
  # Load model
105
  model_path = Path("tasks", "models")
106
+ # model_name = "best.pt" # nano 20e, unpruned
107
+ model_name = "pruned.pt" # nano 20e, 20% pruned
108
  model = YOLO(Path(model_path, model_name))
109
  threshold = 0.14
110