How can I reset the “state value”?

[Atletico1999]

UniPixel is a good piece of work. I'm using it for automated data annotation.

I noticed UniPixel introduced a module called “Object Memory Bank” (perhaps referred to as “state values” or “memory values”) to enable consistent target tracking and inference across frames in video segmentation tasks.

I've customized a class to load the model and process data. I believe this “state value” should be reset to zero when processing new videos to prevent cross-video interference in processing results.

How can I reset this “state value”? I haven't found any corresponding parameters for this module.

Looking forward to your reply! Thank you!

[yeliudev]

Hi @Atletico1999, thanks for your interest in our work!

In our implementation, we do not explicitly save the "states" inside the model. The states are essentially the pred_mask here. So it should be fine if you simply call the model multiple times for different videos. The results shall not be affected.

[Atletico1999]

Thanks for your reply!

Could the "Object Memory Bank" module be reset?

[yeliudev]

I think there is no need to "reset" the module. Could you share your inference pipeline so that I can see fit our model to your use case?

[Atletico1999]

My processing workflow involves passing the frame data corresponding to each video to UniPixel to obtain the mask for each image.

import torch
import numpy as np
import traceback
from typing import List, Optional
from unipixel.dataset.utils import process_vision_info
from unipixel.model.builder import build_model
from unipixel.utils.io import load_image
from unipixel.utils.transforms import get_sam2_transform



class UniPixelWrapper:
    def __init__(self, model_path, device='cuda', dtype='bfloat16'):
        print(f"[Init] Loading UniPixel on {device}...")

        self.model, self.processor = build_model(model_path, device=device, dtype=dtype)
        self.device = next(self.model.parameters()).device

        self.sam2_transform = get_sam2_transform(self.model.config.sam2_image_size)
        self.dtype = self.model.sam2.dtype

        self.reset_state()


    def reset_state(self):
        pass
​

​    @torch.no_grad()
    def predict(self, image_paths: List[str], prompt: str, clear_state: bool = False) -> List[Optional[np.ndarray]]:
        try:
            self.model.seg = []
            if clear_state:
                self.reset_state()

            # 1. loading and stacking
            frames_list = []
            valid_paths = []

            for path in image_paths:
                try:
                    img_tensor = load_image(path)  # (1, H, W, 3)
                    frames_list.append(img_tensor)
                    valid_paths.append(path)
                except Exception as e:
                    print(f"[Load Error] Skipping {path}: {e}")

            if not frames_list:
                return []

            frames = torch.cat(frames_list, dim=0)  # (N, H, W, 3)

            # 2. Construct Prompt (Enable video mode)
            messages = [{
                'role': 'user',
                'content': [{
                    'type': 'video',
                    'video': valid_paths,
                    'min_pixels': 128 * 28 * 28,
                    'max_pixels': 256 * 28 * 28,
                }, {
                    'type': 'text',
                    'text': prompt
                }]
            }]

            # 3. Preprocessing
            text = self.processor.apply_chat_template(messages, add_generation_prompt=True)
            images_processed, videos, kwargs = process_vision_info(messages, return_video_kwargs=True)
            data = self.processor(text=[text], images=images_processed, videos=videos, return_tensors='pt', **kwargs)

            transformed_frames = self.sam2_transform(frames).to(self.dtype)
            data['frames'] = [transformed_frames]
            height, width = frames.shape[1], frames.shape[2]
            data['frame_size'] = [(height, width)]

            # 4. Inference
            output_ids = self.model.generate(
                **data.to(self.device),
                do_sample=False,
                temperature=None,
                top_k=None,
                top_p=None,
                repetition_penalty=None,
                max_new_tokens=512
            )

            # 5. Parse results
            results = []
            seg_output = self.model.seg

            if not seg_output:
                return [np.zeros((height, width), dtype=np.uint8) for _ in valid_paths]

            masks = seg_output[0].squeeze(0).cpu().numpy()  # (N, H, W)
            for i in range(masks.shape[0]):
                mask = (masks[i] * 255).astype(np.uint8)
                results.append(mask)

            self.model.seg = []
            return results

        except Exception as e:
            print(f"[Error] {e}")
            traceback.print_exc()
            return [None] * len(image_paths)



if __name__ == "__main__":
    model = UniPixelWrapper(model_path, device=device)
    for video_frame_paths in videos:
        masks = model.predict(video_frame_paths, prompt, clear_state=True)  # Reset model state values (e.g., Object Memory Bank) before processing frame data for each video

[yeliudev]

In this case, it seems that no "state" is used. So there's no need to reset anything.

[Atletico1999]

Thanks for your reply!