From fc0968b0055553aa24cb200f7bea1339875e7d43 Mon Sep 17 00:00:00 2001 From: Chris Rosewarne Date: Fri, 18 Sep 2026 14:29:52 +1000 Subject: [PATCH 1/2] [feat] Prompt info SEI support --- cfgs/pipeline/split_inference.yaml | 4 +++ compressai_vision/pipelines/base.py | 18 ++++++++++ .../split_inference/image_split_inference.py | 9 +++++ .../split_inference/video_split_inference.py | 36 ++++++++++++++----- 4 files changed, 58 insertions(+), 9 deletions(-) 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..cf73db8 100755 --- a/compressai_vision/pipelines/base.py +++ b/compressai_vision/pipelines/base.py @@ -37,6 +37,7 @@ from pathlib import Path from typing import Callable, Dict, Tuple from uuid import uuid4 as uuid +import yaml import torch import torch.nn as nn @@ -323,6 +324,23 @@ 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..13809da 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,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"]) + # 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 +443,25 @@ 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 ( From 27e0f8d6b99ae1f1252b45ce74b9a74dd0229a67 Mon Sep 17 00:00:00 2001 From: Fabien Racape Date: Wed, 28 Oct 2026 15:55:25 -0700 Subject: [PATCH 2/2] formatting: fix prompt info SEI support --- compressai_vision/pipelines/base.py | 6 ++++-- .../pipelines/split_inference/video_split_inference.py | 6 ++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/compressai_vision/pipelines/base.py b/compressai_vision/pipelines/base.py index cf73db8..f7fd0fd 100755 --- a/compressai_vision/pipelines/base.py +++ b/compressai_vision/pipelines/base.py @@ -37,10 +37,10 @@ from pathlib import Path from typing import Callable, Dict, Tuple from uuid import uuid4 as uuid -import yaml import torch import torch.nn as nn +import yaml from omegaconf.errors import InterpolationResolutionError from torch import Tensor @@ -326,7 +326,9 @@ def _save_hashes(self, dec_hashes): 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}" + 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) diff --git a/compressai_vision/pipelines/split_inference/video_split_inference.py b/compressai_vision/pipelines/split_inference/video_split_inference.py index 13809da..fadc581 100644 --- a/compressai_vision/pipelines/split_inference/video_split_inference.py +++ b/compressai_vision/pipelines/split_inference/video_split_inference.py @@ -334,10 +334,7 @@ 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 - ): + 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 @@ -461,6 +458,7 @@ def stringify(k, v): 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)