diff --git a/cfgs/pipeline/split_inference.yaml b/cfgs/pipeline/split_inference.yaml index a1c4b5e..d9a026e 100755 --- a/cfgs/pipeline/split_inference.yaml +++ b/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" diff --git a/compressai_vision/pipelines/base.py b/compressai_vision/pipelines/base.py index 380a124..f7fd0fd 100755 --- a/compressai_vision/pipelines/base.py +++ b/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, diff --git a/compressai_vision/pipelines/split_inference/image_split_inference.py b/compressai_vision/pipelines/split_inference/image_split_inference.py index c1c9e7e..c184be9 100644 --- a/compressai_vision/pipelines/split_inference/image_split_inference.py +++ b/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: diff --git a/compressai_vision/pipelines/split_inference/video_split_inference.py b/compressai_vision/pipelines/split_inference/video_split_inference.py index 58aa576..fadc581 100644 --- a/compressai_vision/pipelines/split_inference/video_split_inference.py +++ b/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 (