kiurtis commited on
Commit
a6c94b5
·
1 Parent(s): 08ec323
Files changed (1) hide show
  1. tasks/image.py +3 -3
tasks/image.py CHANGED
@@ -37,7 +37,7 @@ print("CUDA version:", torch.version.cuda)
37
  print("cuDNN version:", torch.backends.cudnn.version())
38
 
39
  #MODEL_TYPE = "YOLOv11n"
40
- DESCRIPTION = f"YOLOv11n 1280 not quantisized model with batch 32 inference on TensorRT"
41
  print(14)
42
  ROUTE = "/image"
43
  print(15)
@@ -141,11 +141,11 @@ async def evaluate_image(request: ImageEvaluationRequest):
141
  #PATH_TO_MODEL = 'models/best_YOLOv11n_1280.onnx'
142
  #PATH_TO_MODEL = 'models/best_yolov6n_1280.pt'
143
  #PATH_TO_MODEL = 'models/best_YOLOv11n_1280_real_half.onnx'
144
- PATH_TO_MODEL = 'models/best_YOLOv11n_1280_half_batch_32.engine'
145
  INFERENCE_ENGINE_TYPE = 'pt'
146
  INPUT_SIZE = 1280
147
  N_TEST_BATCHES = 2
148
- BATCH_SIZE = 32 # Can be adjusted as needed
149
  print("PATH_TO_MODEL", PATH_TO_MODEL)
150
  def preprocessor(frame):
151
  #frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Only when read from file
 
37
  print("cuDNN version:", torch.backends.cudnn.version())
38
 
39
  #MODEL_TYPE = "YOLOv11n"
40
+ DESCRIPTION = f"YOLOv11n 1280 not quantisized model with batch 1 inference on TensorRT"
41
  print(14)
42
  ROUTE = "/image"
43
  print(15)
 
141
  #PATH_TO_MODEL = 'models/best_YOLOv11n_1280.onnx'
142
  #PATH_TO_MODEL = 'models/best_yolov6n_1280.pt'
143
  #PATH_TO_MODEL = 'models/best_YOLOv11n_1280_real_half.onnx'
144
+ PATH_TO_MODEL = 'models/best_YOLOv11n_1280_full_batch_1.engine'
145
  INFERENCE_ENGINE_TYPE = 'pt'
146
  INPUT_SIZE = 1280
147
  N_TEST_BATCHES = 2
148
+ BATCH_SIZE = 1 # Can be adjusted as needed
149
  print("PATH_TO_MODEL", PATH_TO_MODEL)
150
  def preprocessor(frame):
151
  #frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Only when read from file