Simplify model names and sizes

README.md

@@ -44,11 +44,11 @@ Here are models that can be downloaded:
 
 | Model             | Parameters | Size  | Download                     |
 |-------------------|------------|-------|------------------------------|
-| SAM 91M           | 91M        | 100MB | `osam run sam:91m`           |
-| SAM 308M          | 308M       | 320MB | `osam run sam:308m`          |
-| SAM 636M          | 636M       | 630MB | `osam run sam`               |
-| EfficientSAM 10M | 10M        | 40MB  | `osam run efficientsam:10m` |
-| EfficientSAM 25M | 25M        | 100MB | `osam run efficientsam`     |
+| SAM 100M          | 100M       | 100MB | `osam run sam:100m`           |
+| SAM 300M          | 300M       | 300MB | `osam run sam:300m`          |
+| SAM 600M          | 600M       | 600MB | `osam run sam`               |
+| EfficientSAM 10M  | 10M        | 40MB  | `osam run efficientsam:10m`  |
+| EfficientSAM 30M  | 30M        | 100MB | `osam run efficientsam`      |
 
 PS. `sam`, `efficientsam` is equivalent to `sam:latest`, `efficientsam:latest`.
 

scripts/count_parameter_size.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+
+import argparse
+
+from loguru import logger
+import onnx
+
+
+def count_parameters(onnx_file):
+    model = onnx.load(onnx_file)
+    total_params = 0
+
+    for initializer in model.graph.initializer:
+        param_count = 1
+        for dim in initializer.dims:
+            param_count *= dim
+        total_params += param_count
+
+    return total_params
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("onnx_file")
+    args = parser.parse_args()
+
+    logger.debug(f"ONNX file: {args.onnx_file!r}")
+    num_params = count_parameters(onnx_file=args.onnx_file)
+    logger.debug(f"Number of parameters: {num_params} = {num_params / 1e6:.2f}M = {num_params / 1e9:.2f}B")
+
+
+if __name__ == "__main__":
+    main()