Multiple updates and fixes for video

add tests for codecs

remove collect.py, can be done using bench when per sequence json files are already present in output dir

fix default metric and check if output provided

align eval metrics with bench

update tests for new eval/bench

fix test bench video

add eval for video with ANS

edit output json files

rename eval output files  and fix  test

fix update() for ssf2020

minor cleaning, add model info in per sequence results

add pretrained model quality level in output results

actually write bitstreams and check size for bitrate calculation

fix video codec

add possibility to keep bitstreams

eval with zero padding and getting reconstructed from encode

fix ssf2020 decompress, are compress/decompress needed in video (unused in codec.py and eval)

minor fixes and pytests

update pre-trained models

.gitlab-ci.yml

@@ -73,6 +73,7 @@ isort:
     - docker
 
 test:
+  cache: []
   stage: test
   image: pytorch/pytorch:$PYTORCH_IMAGE
   before_script:

compressai/models/video/google.py

@@ -429,7 +429,7 @@ def decompress(self, strings, shapes):
         for i in range(1, len(strings)):
             string = strings[i]
             shape = shapes[i]
-            x_ref, out_interframe = self.decoder_inter(x_ref, string, shape)
+            x_ref = self.decode_inter(x_ref, string, shape)
             dec_frames.append(x_ref)
 
         return dec_frames
@@ -492,16 +492,17 @@ def update(self, scale_table=None, force=False):
         updated = self.img_hyperprior.gaussian_conditional.update_scale_table(
             scale_table, force=force
         )
-        # updated |= super().update(force=force)
 
-        updated = self.res_hyperprior.gaussian_conditional.update_scale_table(
+        updated |= self.img_hyperprior.entropy_bottleneck.update(force=force)
+
+        updated |= self.res_hyperprior.gaussian_conditional.update_scale_table(
             scale_table, force=force
         )
-        # updated |= super().update(force=force)
+        updated |= self.res_hyperprior.entropy_bottleneck.update(force=force)
 
-        updated = self.motion_hyperprior.gaussian_conditional.update_scale_table(
+        updated |= self.motion_hyperprior.gaussian_conditional.update_scale_table(
             scale_table, force=force
         )
-        # updated |= super().update(force=force)
+        updated |= self.motion_hyperprior.entropy_bottleneck.update(force=force)
 
         return updated

compressai/utils/eval_model/__main__.py

@@ -49,8 +49,8 @@
 
 import compressai
 
+from compressai.zoo import image_models as pretrained_models
 from compressai.zoo import load_state_dict
-from compressai.zoo import models as pretrained_models
 from compressai.zoo.image import model_architectures as architectures
 
 torch.backends.cudnn.deterministic = True

compressai/utils/video/bench/__main__.py

@@ -80,23 +80,22 @@ def func(codec, i, filepath, qp, outputdir, cuda, force, dry_run):
         with sequence_metrics_path.open("r") as f:
             metrics = json.load(f)["results"]
         return i, qp, metrics
-
-    with tempfile.NamedTemporaryFile(suffix=".yuv", delete=True) as f:
-        # decode sequence
-        decode_cmd = codec.get_decode_cmd(binpath, f.name, filepath)
-        run_cmdline(decode_cmd)
-
-        # compute metrics
-        metrics = evaluate(filepath, Path(f.name), binpath, cuda)
-        output = {
-            "source": filepath.stem,
-            "name": codec.name_config(),
-            "description": codec.description(),
-            "results": metrics,
-        }
-        with sequence_metrics_path.open("wb") as f:
-            f.write(json.dumps(output, indent=2).encode())
-    print(type(metrics))
+    else:
+        with tempfile.NamedTemporaryFile(suffix=".yuv", delete=True) as f:
+            # decode sequence
+            decode_cmd = codec.get_decode_cmd(binpath, f.name, filepath)
+            run_cmdline(decode_cmd)
+
+            # compute metrics
+            metrics = evaluate(filepath, Path(f.name), binpath, cuda)
+            output = {
+                "source": filepath.stem,
+                "name": codec.name_config(),
+                "description": codec.description(),
+                "results": metrics,
+            }
+            with sequence_metrics_path.open("wb") as f:
+                f.write(json.dumps(output, indent=2).encode())
     return i, qp, metrics
 
 
@@ -205,7 +204,9 @@ def evaluate(
         k: torch.mean(torch.stack(v)) for k, v in results.items()
     }
     filesize = get_filesize(bitstream_path)
-    seq_results["bitrate"] = float(filesize * org_seq.framerate / (num_frames * 1000))
+    seq_results["bitrate"] = float(
+        filesize * 8 * org_seq.framerate / (num_frames * 1000)
+    )
 
     seq_results["psnr-rgb"] = (
         20 * np.log10(max_val) - 10 * torch.log10(seq_results.pop("mse-rgb")).item()
@@ -224,21 +225,6 @@ def evaluate(
     return seq_results
 
 
-def aggregate_results(filepaths: List[Path]) -> Dict[str, Any]:
-    metrics = defaultdict(list)
-
-    # sum
-    for f in filepaths:
-        with f.open("r") as fd:
-            data = json.load(fd)
-        for k, v in data.items():
-            metrics[k].append(v)
-
-    # normalize
-    agg = {k: np.mean(v) for k, v in metrics.items()}
-    return agg
-
-
 def collect(
     dataset: Path,
     codec_class: Codec,

compressai/utils/video/bench/codecs.py

@@ -102,7 +102,7 @@ def name(self):
         return "x264"
 
     def description(self):
-        return f"{self.name} {self.preset}, tuned for {self.tune}, ffmpeg version {get_ffmpeg_version()}"
+        return f"{self.name} {self.preset}, {self.tune}, ffmpeg version {get_ffmpeg_version()}"
 
     def name_config(self):
         return f"{self.name}-{self.preset}-tune-{self.tune}"
@@ -131,6 +131,7 @@ def get_encode_cmd(self, filepath: Path, qp, bindir) -> List[Any]:
         binpath = self.get_bin_path(filepath, qp, bindir)
         cmd = [
             "ffmpeg",
+            "-y",
             "-s:v",
             f"{info['width']}x{info['height']}",
             "-i",