Update model (#210)

* use 640 image size

* reduce detection rate

* reduce detection rate

* adpt test size

* update model sha

pyroengine/vision.py

@@ -41,7 +41,7 @@ class Classifier:
         model_path: model path
     """
 
-    def __init__(self, model_path: Optional[str] = "data/model.onnx", base_img_size: int = 1024) -> None:
+    def __init__(self, model_path: Optional[str] = "data/model.onnx", base_img_size: int = 640) -> None:
         if model_path is None:
             model_path = "data/model.onnx"
 

src/run.py

@@ -92,7 +92,7 @@ def main(args):
     )
     # Model
     parser.add_argument("--model_path", type=str, default="data/model.onnx", help="model path")
-    parser.add_argument("--thresh", type=float, default=0.25, help="Confidence threshold")
+    parser.add_argument("--thresh", type=float, default=0.15, help="Confidence threshold")
     # Camera & cache
     parser.add_argument("--creds", type=str, default="data/credentials.json", help="Camera credentials")
     parser.add_argument("--cache", type=str, default="./data", help="Cache folder")

tests/test_vision.py

@@ -8,7 +8,7 @@
 
 METADATA_PATH = "data/model_metadata.json"
 model_path = "data/model.onnx"
-sha = "12b9b5728dfa2e60502dcde2914bfdc4e9378caa57611c567a44cdd6228838c2"
+sha = "9f1b1c2654d98bbed91e514ce20ea73a0a5fbd1111880f230d516ed40ea2dc58"
 
 
 def custom_isfile_false(path):
@@ -30,21 +30,21 @@ def test_classifier(mock_wildfire_image):
         # Instantiate the ONNX model
         model = Classifier()
         # Check preprocessing
-        out = model.preprocess_image(mock_wildfire_image, (1024, 576))
+        out = model.preprocess_image(mock_wildfire_image, (640, 384))
         assert isinstance(out, np.ndarray) and out.dtype == np.float32
-        assert out.shape == (1, 3, 576, 1024)
+        assert out.shape == (1, 3, 384, 640)
         # Check inference
         out = model(mock_wildfire_image)
         assert out.shape == (1, 5)
         conf = np.max(out[:, 4])
         assert conf >= 0 and conf <= 1
 
         # Test mask
-        mask = np.ones((1024, 576))
+        mask = np.ones((640, 384))
         out = model(mock_wildfire_image, mask)
         assert out.shape == (1, 5)
 
-        mask = np.zeros((1024, 1024))
+        mask = np.zeros((640, 384))
         out = model(mock_wildfire_image, mask)
         assert out.shape == (0, 5)
         os.remove(model_path)