feat: refactor to use only singular classifcation model for color (#533)

* feat: refactor to use only singular classifcation model for color

Author: itisacloud

sketch_map_tool/config.py

@@ -22,8 +22,7 @@
     "wms-read-timeout": 600,
     "max-nr-simultaneous-uploads": 100,
     "max_pixel_per_image": 10e8,  # 10.000*10.000
-    "yolo_osm_cls": "SMT-OSM-CLS",
-    "yolo_esri_cls": "SMT-ESRI-CLS",
+    "yolo_cls": "SMT-CLS",
     "yolo_osm_obj": "SMT-OSM",
     "yolo_esri_obj": "SMT-ESRI",
     "model_type_sam": "vit_b",

sketch_map_tool/routes.py

@@ -99,14 +99,6 @@ def weights_smt_osm(lang="en") -> Response:  # pyright: ignore
     return send_from_directory(dir, name, as_attachment=True)
 
 
-@app.get("/weights/SMT-OSM-CLS.pt")
-@app.get("/<lang>/weights/SMT-OSM-CLS.pt")
-def weights_smt_osm_cls(lang="en") -> Response:  # pyright: ignore
-    dir = Path(config.get_config_value("weights-dir"))
-    name = "SMT-OSM-CLS.pt"
-    return send_from_directory(dir, name, as_attachment=True)
-
-
 @app.get("/weights/SMT-ESRI.pt")
 @app.get("/<lang>/weights/SMT-ESRI.pt")
 def weights_smt_esri(lang="en") -> Response:  # pyright: ignore
@@ -115,11 +107,11 @@ def weights_smt_esri(lang="en") -> Response:  # pyright: ignore
     return send_from_directory(dir, name, as_attachment=True)
 
 
-@app.get("/weights/SMT-ESRI-CLS.pt")
-@app.get("/<lang>/weights/SMT-ESRI-CLS.pt")
-def weights_smt_esri_cls(lang="en") -> Response:  # pyright: ignore
+@app.get("/weights/SMT-CLS.pt")
+@app.get("/<lang>/weights/SMT-CLS.pt")
+def weights_smt_cls(lang="en") -> Response:  # pyright: ignore
     dir = Path(config.get_config_value("weights-dir"))
-    name = "SMT-ESRI-CLS.pt"
+    name = "SMT-CLS.pt"
     return send_from_directory(dir, name, as_attachment=True)
 
 

sketch_map_tool/tasks.py

@@ -51,9 +51,8 @@ def init_worker_ml_models(**_):
     logging.info("Initialize ml-models.")
     global sam_predictor
     global yolo_obj_osm
-    global yolo_cls_osm
     global yolo_obj_esri
-    global yolo_cls_esri
+    global yolo_cls
 
     path = init_sam2()
     device = select_computation_device()
@@ -65,9 +64,8 @@ def init_worker_ml_models(**_):
     sam_predictor = SAM2ImagePredictor(sam2_model)
 
     yolo_obj_osm = YOLO_MB(init_model(get_config_value("yolo_osm_obj")))
-    yolo_cls_osm = YOLO(init_model(get_config_value("yolo_osm_cls")))
     yolo_obj_esri = YOLO_MB(init_model(get_config_value("yolo_esri_obj")))
-    yolo_cls_esri = YOLO(init_model(get_config_value("yolo_esri_cls")))
+    yolo_cls = YOLO(init_model(get_config_value("yolo_cls")))
 
 
 @worker_process_shutdown.connect
@@ -140,10 +138,8 @@ def digitize_sketches(
 ) -> FeatureCollection:
     if layer == "osm":
         yolo_obj = yolo_obj_osm
-        yolo_cls = yolo_cls_osm
     elif layer == "esri-world-imagery":
         yolo_obj = yolo_obj_esri
-        yolo_cls = yolo_cls_esri
     else:
         raise ValueError("Unexpected layer: " + layer)
 

tests/integration/upload_processing/test_detect_markings.py

@@ -41,13 +41,6 @@ def yolo_osm_obj() -> YOLO_MB:
     return YOLO_MB(path)
 
 
-@pytest.fixture
-def yolo_osm_cls() -> YOLO:
-    """YOLO Classification"""
-    path = init_model(get_config_value("yolo_osm_cls"))
-    return YOLO(path)
-
-
 @pytest.fixture
 def yolo_esri_obj() -> YOLO_MB:
     """YOLO Object Detection"""
@@ -56,9 +49,9 @@ def yolo_esri_obj() -> YOLO_MB:
 
 
 @pytest.fixture
-def yolo_esri_cls() -> YOLO:
+def yolo_cls() -> YOLO:
     """YOLO Classification"""
-    path = init_model(get_config_value("yolo_osm_cls"))
+    path = init_model(get_config_value("yolo_cls"))
     return YOLO(path)
 
 
@@ -68,17 +61,14 @@ def test_detect_markings(
     map_frame_marked,
     map_frame,
     yolo_osm_obj,
-    yolo_osm_cls,
     yolo_esri_obj,
-    yolo_esri_cls,
+    yolo_cls,
     sam_predictor,
 ):
     if layer.value == "osm":
         yolo_obj = yolo_osm_obj
-        yolo_cls = yolo_osm_cls
     else:
         yolo_obj = yolo_esri_obj
-        yolo_cls = yolo_esri_cls
     markings = detect_markings(
         map_frame_marked,
         np.asarray(Image.open(map_frame)),

tests/integration/upload_processing/test_ml_models.py

@@ -9,10 +9,9 @@
 @pytest.mark.parametrize(
     "id",
     (
-        get_config_value("yolo_osm_cls"),
-        get_config_value("yolo_esri_cls"),
         get_config_value("yolo_osm_obj"),
         get_config_value("yolo_esri_obj"),
+        get_config_value("yolo_cls"),
     ),
 )
 def test_init_model(id):

tests/unit/test_config.py

@@ -34,8 +34,7 @@ def config_keys():
         "wms-read-timeout",
         "max-nr-simultaneous-uploads",
         "max_pixel_per_image",
-        "yolo_osm_cls",
-        "yolo_esri_cls",
+        "yolo_cls",
         "yolo_osm_obj",
         "yolo_esri_obj",
         "model_type_sam",