fix: pytest for evaluating model performance dont read stderr anymore

Author: fracape

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",
@@ -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,7 @@ 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

@@ -78,7 +78,7 @@ def test_eval_model():
 @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
+    model, quality, metric, entropy_estimation, tmpdir
 ):
     here = os.path.dirname(__file__)
     dirpath = os.path.join(here, "assets/dataset/image")
@@ -92,13 +92,17 @@ def test_eval_model_pretrained(
         metric,
         "-q",
         quality,
+        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",

tests/test_eval_model_video.py

@@ -80,7 +80,7 @@ def test_eval_model_video():
 @pytest.mark.parametrize("metric", ("mse",))
 @pytest.mark.parametrize("entropy_estimation", (True, False))
 def test_eval_model_pretrained(
-    capsys, model, quality, metric, entropy_estimation, tmpdir
+    model, quality, metric, entropy_estimation, tmpdir
 ):
     here = os.path.dirname(__file__)
     dirpath = os.path.join(here, "assets/dataset/video")
@@ -95,13 +95,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",