[feat] Prompt info SEI support

cfgs/pipeline/split_inference.yaml

@@ -38,6 +38,10 @@ codec:
     load_hashes: False
     hash_dir_enc: "${codec.output_dir}/features_pre_nn_part2"
     hash_dir_dec: "${codec.hash_dir}/hashes_dec"
+    load_prompts: False
+    save_prompts: False
+    prompt_file_enc: "prompt_file_enc.yaml"
+    prompt_file_dec: "propmt_file_dec.yaml"
 nn_task_part2:
     dump_results: False
     output_results_dir: "${codec.output_dir}/output_results"

compressai_vision/pipelines/base.py

@@ -40,6 +40,7 @@
 
 import torch
 import torch.nn as nn
+import yaml
 
 from omegaconf.errors import InterpolationResolutionError
 from torch import Tensor
@@ -323,6 +324,25 @@ def _save_hashes(self, dec_hashes):
                 for dec_hash in dec_hashs:
                     f.write(dec_hash.hex() + "\n")
 
+    def _load_prompts(self):
+        prompt_file_enc = Path(self.configs["codec"].prompt_file_enc)
+        assert (
+            prompt_file_enc.is_file()
+        ), f"Prompt info SEI message input file not found: {prompt_file_enc}"
+
+        with open(prompt_file_enc) as f:
+            prompts = yaml.safe_load(f)
+
+        return prompts
+
+    def _save_prompts(self, dec_prompts):
+        prompt_file_dec = Path(self.configs["codec"].prompt_file_dec)
+        if prompt_file_dec.is_file():
+            self.logger.warning(f"Overwriting prompt file: {prompt_file_dec}")
+
+        with open(prompt_file_dec, "w") as f:
+            yaml.dump(prompt_file_dec, f)
+
     def _from_input_to_features(
         self,
         vision_model: BaseWrapper,

compressai_vision/pipelines/split_inference/image_split_inference.py

@@ -136,6 +136,9 @@ def __call__(
                 if self.configs["codec"]["load_hashes"]:
                     features["hashes"] = self._load_hashes()
 
+                if self.configs["codec"]["load_prompts"]:
+                    features["prompt_info"] = self._load_prompts()
+
                 del res["data"]
 
                 start = time_measure()
@@ -214,6 +217,12 @@ def __call__(
             ):
                 self._save_hashes(dec_features["decoded_feature_tensor_hash"])
 
+            if (
+                "prompt_info" in dec_features
+                and dec_features["prompt_info"] is not None
+            ):
+                self._save_prompts(dec_features["prompt_info"])
+
             # dec_features should contain "org_input_size" and "input_size"
             # When using anchor codecs, that's not the case, we read input images to derive them
             if "vision_model_info" not in dec_features:

compressai_vision/pipelines/split_inference/video_split_inference.py

@@ -242,6 +242,9 @@ def __call__(
             if self.configs["codec"]["load_hashes"]:
                 features["hashes"] = self._load_hashes()
 
+            if self.configs["codec"]["load_prompts"]:
+                features["prompt_info"] = self._load_prompts()
+
             if not evaluator.calculate_feature_mse:
                 self._input_ftensor_buffer = []
 
@@ -331,6 +334,9 @@ def __call__(
         ):
             self._save_hashes(dec_features["decoded_feature_tensor_hash"])
 
+        if "prompt_info" in dec_features and dec_features["prompt_info"] is not None:
+            self._save_prompts(dec_features["prompt_info"])
+
         # separate a tensor of each keyword item into a list of tensors
         dec_ftensors_list = self._feature_tensor_dict_to_list(dec_features["data"])
         assert all(
@@ -434,16 +440,26 @@ def __call__(
             out_res["org_input_size"] = (
                 f"{dec_features['org_input_size']['height']}x{dec_features['org_input_size']['width']}"
             )
-            if (
-                "decoded_feature_tensor_hash" in dec_features
-                and dec_features["decoded_feature_tensor_hash"] is not None
-            ):
-                out_res["decoded_feature_tensor_hash"] = (
-                    dec_features["decoded_feature_tensor_hash"][e]
-                    if e in dec_features["decoded_feature_tensor_hash"]
-                    else None
-                )
+            # Convert frame dict to framewise list
+            for framewise_sei in ["prompt_info", "decoded_feature_tensor_hash"]:
+                if (
+                    framewise_sei in dec_features
+                    and dec_features[framewise_sei] is not None
+                ):
+                    out_res[framewise_sei] = (
+                        dec_features[framewise_sei][e]
+                        if e in dec_features[framewise_sei]
+                        else None
+                    )
+
+            # Stringify SEI output
+            def stringify(k, v):
+                if k in ["prompt_info"]:
+                    print(v)
+                    return str(v)
+                return v
 
+            out_res = {k: stringify(k, v) for k, v in out_res.items()}
             output_list.append(out_res)
 
         if (