fix: pytest for evaluating model performance dont read stderr anymore

Author: fracape

.github/workflows/pytest.yml

@@ -11,16 +11,19 @@ jobs:
           - "3.8"
           - "3.9"
           - "3.10"
+          - "3.11"
+          - "3.12"
 
     steps:
     - uses: actions/checkout@v4
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Install Python dependencies
-      run: |
-        python -m pip install -U pip
-        pip install -e .[test]
+
+    - name: Install the project
+        run: uv sync --locked --all-extras --dev
+
     - name: Run unit tests
-      run: python -m pytest -m "not slow" --cov=compressai -s tests/
+      run: uv run pytest -m "not slow" --cov=compressai -s tests/

.github/workflows/static-analysis.yml

@@ -10,17 +10,10 @@ jobs:
         python-version:
           - "3.8"
           - "3.10"
+          - "3.12"
+
         include:
           - os: "ubuntu-latest"
     steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install Python dependencies
-      run: |
-        python3 -m pip install -U pip
-        python3 -m pip install .[dev]
-    - name: Run static analysis checks
+    - uses: astral-sh/ruff-action@v3
       run: make static-analysis

compressai/entropy_models/entropy_models.py

@@ -847,9 +847,9 @@ def _build_cdf(self, scales, means, weights, abs_max):
             ),
             dim=1,
         ).transpose(0, 1)
-        pmf_quantized[pmf_real_steal_indices[0], pmf_real_steal_indices[1]] -= (
-            pmf_zero_count
-        )
+        pmf_quantized[
+            pmf_real_steal_indices[0], pmf_real_steal_indices[1]
+        ] -= pmf_zero_count
 
         cdf = F.pad(torch.cumsum(pmf_quantized, 1).int(), (1, 0), "constant", 0)
         cdf = F.pad(cdf, (0, 1), "constant", cdf_limit + 1)

compressai/utils/video/eval_model/__main__.py

@@ -467,6 +467,13 @@ def create_parser() -> argparse.ArgumentParser:
         default="mse",
         help="metric trained against (default: %(default)s)",
     )
+    parent_parser.add_argument(
+        "-d",
+        "--output_directory",
+        type=str,
+        default="",
+        help="path of output directory. Optional, required for output json file, results per video.",
+    )
     parent_parser.add_argument(
         "-o",
         "--output-file",
@@ -506,7 +513,7 @@ def create_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def main(args: Any = None) -> None:
+def main(args: Any = None) -> None:  # noqa: C901
     if args is None:
         args = sys.argv[1:]
     parser = create_parser()
@@ -525,8 +532,8 @@ def main(args: Any = None) -> None:
         raise SystemExit(1)
 
     # create output directory
-    outputdir = args.output
-    Path(outputdir).mkdir(parents=True, exist_ok=True)
+    if args.output_directory:
+        Path(args.output_directory).mkdir(parents=True, exist_ok=True)
 
     if args.source == "pretrained":
         args.qualities = [int(q) for q in args.qualities.split(",") if q]
@@ -561,7 +568,7 @@ def main(args: Any = None) -> None:
             filepaths,
             args.dataset,
             model,
-            outputdir,
+            args.output_directory,
             trained_net=trained_net,
             description=description,
             **args_dict,
@@ -581,7 +588,9 @@ def main(args: Any = None) -> None:
     else:
         output_file = args.output_file
 
-    with (Path(f"{outputdir}/{output_file}").with_suffix(".json")).open("wb") as f:
+    with (Path(f"{args.output_directory}/{output_file}").with_suffix(".json")).open(
+        "wb"
+    ) as f:
         f.write(json.dumps(output, indent=2).encode())
     print(json.dumps(output, indent=2))
 

tests/test_eval_model.py

@@ -77,9 +77,7 @@ def test_eval_model():
 @pytest.mark.parametrize("quality", ("1", "4", "8"))
 @pytest.mark.parametrize("metric", ("mse", "ms-ssim"))
 @pytest.mark.parametrize("entropy_estimation", (False, True))
-def test_eval_model_pretrained(
-    capsys, model, quality, metric, entropy_estimation, tmpdir
-):
+def test_eval_model_pretrained(model, quality, metric, entropy_estimation, tmpdir):
     here = os.path.dirname(__file__)
     dirpath = os.path.join(here, "assets/dataset/image")
 
@@ -92,13 +90,18 @@ def test_eval_model_pretrained(
         metric,
         "-q",
         quality,
+        "-o",
+        f"{model}-{metric}-{quality}",
+        "-d",
+        str(tmpdir),
     ]
     if entropy_estimation:
         cmd += ["--entropy-estimation"]
     eval_model.main(cmd)
 
-    output = capsys.readouterr().out
-    output = json.loads(output)
+    with open(f"{tmpdir}/{model}-{metric}-{quality}.json") as f:
+        output = json.load(f)
+
     expected = os.path.join(
         here,
         "expected",

tests/test_eval_model_video.py

@@ -79,9 +79,7 @@ def test_eval_model_video():
 @pytest.mark.parametrize("quality", ("1", "4", "8"))
 @pytest.mark.parametrize("metric", ("mse",))
 @pytest.mark.parametrize("entropy_estimation", (True, False))
-def test_eval_model_pretrained(
-    capsys, model, quality, metric, entropy_estimation, tmpdir
-):
+def test_eval_model_pretrained(model, quality, metric, entropy_estimation, tmpdir):
     here = os.path.dirname(__file__)
     dirpath = os.path.join(here, "assets/dataset/video")
 
@@ -95,13 +93,18 @@ def test_eval_model_pretrained(
         metric,
         "-q",
         quality,
+        "-d",
+        str(tmpdir),
+        "-o",
+        f"{model}-{metric}-{quality}",
     ]
     if entropy_estimation:
         cmd += ["--entropy-estimation"]
     eval_model.main(cmd)
 
-    output = capsys.readouterr().out
-    output = json.loads(output)
+    with open(f"{tmpdir}/{model}-{metric}-{quality}.json") as f:
+        output = json.load(f)
+
     expected = os.path.join(
         here,
         "expected",