dev(narugo): use new models

monochrome.py

@@ -17,19 +17,18 @@ __all__ = [
 
 # _DEFAULT_MONOCHROME_CKPT = 'monochrome-resnet18-safe2-450.onnx'
 _MONOCHROME_CKPTS = [
-    'monochrome-caformer-110.onnx',
-    'monochrome-caformer_safe2-80.onnx',
-    'monochrome-caformer_safe4-70.onnx',
-    'monochrome-caformer-40.onnx',
+    'mobilenetv3_large_100_safe2',
+    'mobilenetv3_large_100',
+    'caformer_s36',
 ]
 _DEFAULT_MONOCHROME_CKPT = _MONOCHROME_CKPTS[0]
 
 
 @lru_cache()
-def _monochrome_validate_model(ckpt):
+def _monochrome_validate_model(model):
     return open_onnx_model(hf_hub_download(
-        'deepghs/imgutils-models',
-        f'monochrome/{ckpt}'
+        f'deepghs/monochrome_detect',
+        f'{model}/model.onnx'
     ))
 
 
@@ -49,12 +48,12 @@ def _2d_encode(image: Image.Image, size: Tuple[int, int] = (384, 384),
     return data
 
 
-def get_monochrome_score(image: ImageTyping, ckpt: str = _DEFAULT_MONOCHROME_CKPT) -> float:
+def get_monochrome_score(image: ImageTyping, model: str = _DEFAULT_MONOCHROME_CKPT):
     image = load_image(image, mode='RGB')
     input_data = _2d_encode(image).astype(np.float32)
     input_data = np.stack([input_data])
-    output_data, = _monochrome_validate_model(ckpt).run(['output'], {'input': input_data})
-    return {name: v.item() for name, v in zip(['normal', 'monochrome'], output_data[0])}
+    output_data, = _monochrome_validate_model(model).run(['output'], {'input': input_data})
+    return {name: v.item() for name, v in zip(['monochrome', 'normal'], output_data[0])}
 
 
 def is_monochrome(image: ImageTyping, threshold: float = 0.5, ckpt: str = _DEFAULT_MONOCHROME_CKPT) -> bool: