Does UniPixel support Image Regional Segmentation (e.g., point/box prompts like SAM)?

[lif314]

Hi, thank you for the great work!

I would like to know if this model supports image regional segmentation. Specifically, can it take user-provided prompts such as points or bounding boxes to segment objects, similar to how SAM (Segment Anything Model) works?

For example, given a text instruction like:

"Please segment the object based on the referring points/bbox."

would the model be able to leverage the spatial hints (points or bounding boxes) along with the text to produce the desired segmentation?

If not currently supported, are there any plans to add such functionality in the future?

Thank you!

[lif314]

@yeliudev Hi !

I have written a test script based on my understanding of the PixelQA loading format. However, the results I'm getting are not good.

I'm wondering whether the issue lies with: (1) The prompt design, or (2) My script missing some necessary preprocessing steps (e.g., specific image/video formatting, visual prompt encoding, or mask post-processing)?

I tried to follow the same loading format used in PixelQA as closely as possible, but the performance is significantly worse than expected.

Could you please provide some guidance on what might be going wrong? Any suggestions for debugging or improving the script would be greatly appreciated.

Thank you!

• Input:

• Result

• Test script

# Copyright (c) 2025 Ye Liu. Licensed under the BSD-3-Clause License.

import argparse
import re
import torch.nn.functional as F
import torchvision.transforms.functional as T

import imageio.v3 as iio
import nncore
import torch
import random
import numpy as np
from PIL import Image

from unipixel.dataset.utils import process_vision_info
from unipixel.model.builder import build_model
from unipixel.utils.io import load_image, load_video
from unipixel.utils.transforms import get_sam2_transform
from unipixel.utils.visualizer import draw_mask
from unipixel.constants import MEM_TOKEN, SEG_TOKEN
from sam2.modeling.sam2_utils import get_next_point, sample_box_points

BANNER = r"""
=================================================================================
   __    __    __   __    __    ______     __   ___   ___   _______    __
  |  |  |  |  |  \ |  |  |  |  |   _  \   |  |  \  \ /  /  |   ____|  |  |
  |  |  |  |  |   \|  |  |  |  |  |_)  |  |  |   \  V  /   |  |__     |  |
  |  |  |  |  |  . `  |  |  |  |   ___/   |  |    >   <    |   __|    |  |
  |  `--'  |  |  |\   |  |  |  |  |       |  |   /  .  \   |  |____   |  `----.
   \______/   |__| \__|  |__|  | _|       |__|  /__/ \__\  |_______|  |_______|

=================================================================================
"""

INFO = """
\033[1;36m   Examples:\033[0m 1. Please segment the tallest giraffe.
             2. Where is the nearest sheep? Please provide the segmentation mask.
             3. Why is the boy crying? Please provide the segmentation mask and explain why.
             4. Who shooted the ball? Please answer the question and provide the segmentation mask.
             5. Please segment the object according to the description: <a-long-description>

\033[1;32m Model Path:\033[0m {}
\033[1;32m Media Path:\033[0m {}
\033[1;32m     Prompt:\033[0m {}
"""


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--media_path', default='demo/pcd0900r.png')
    parser.add_argument('--prompt', default='Please locate the single object according to the reference points.')
    parser.add_argument('--refer_path', default='demo/pcd0900r_mask.png')
    parser.add_argument('--output_dir', default='outputs')
    parser.add_argument('--model_path', default='PolyU-ChenLab/UniPixel-3B')
    parser.add_argument('--sample_frames', type=int, default=16)
    parser.add_argument('--device', default='auto')
    parser.add_argument('--dtype', default='bfloat16')
    args = parser.parse_args()
    return args

def load_mask_from_path(mask_path):
    mask_img = Image.open(mask_path).convert('L')

    mask_np = np.array(mask_img) / 255.0
    mask_binary = mask_np > 0.5

    return mask_binary

if __name__ == '__main__':
    args = parse_args()

    print(BANNER + INFO.format(args.model_path, args.media_path, args.prompt))
    prompt_idx = 1
    model, processor = build_model(args.model_path, device=args.device, dtype=args.dtype)
    device = next(model.parameters()).device

    sam2_transform = get_sam2_transform(model.config.sam2_image_size)

    if any(args.media_path.endswith(k) for k in ('jpg', 'png')):
        frames, images = load_image(args.media_path), [args.media_path]
    else:
        frames, images = load_video(args.media_path, sample_frames=args.sample_frames)

    frame_size = frames.shape[1:3]

    video = None
    if video is None:
        prefix = f'Here is an image with the following highlighted regions:\n[0]: <{prompt_idx}> {MEM_TOKEN}\n'
    else:
        prefix = f'Here is a video with {len(images)} frames denoted as <1> to <{len(images)}>. The highlighted regions are as follows:\n[0]: <{prompt_idx}>-<{prompt_idx + 1}> {MEM_TOKEN}\n'

    # Please localize the relevant objects with IDs.

    messages = [{
        'role':
        'user',
        'content': [{
            'type': 'video',
            'video': images,
            'min_pixels': 128 * 28 * 28,
            'max_pixels': 256 * 28 * 28 * int(args.sample_frames / len(images))
        }, {
            'type': 'text',
            'text': args.prompt
        }]
    }]

    text = processor.apply_chat_template(messages, add_generation_prompt=True)

    images, videos, kwargs = process_vision_info(messages, return_video_kwargs=True)

    data = processor(text=[text], images=images, videos=videos, return_tensors='pt', **kwargs)

    masks = []
    inds = [0]
    for raw_obj_masks in [[args.refer_path]]:
        obj_masks = []
        for i in inds:
            label = raw_obj_masks[i]
            if isinstance(label, str):
                label = np.array(Image.open(label))
            elif label is None:
                label = np.full(frame_size, -1)
            obj_masks.append(torch.from_numpy(label == 255).float())
        masks.append(obj_masks)
    label_mask = [masks]

    label_mask = torch.stack([torch.stack(o) for t in label_mask for o in t]).transpose(0, 1)
    assert label_mask.shape[2:] == frame_size
    label_mask = T.resize(label_mask, (model.config.sam2_image_size, model.config.sam2_image_size))
    label_mask = label_mask > 0

    print(f"label_mask shape: {label_mask.shape}")

    if label_mask.size(1) == 1:
        # data_types = [random.choice(['point', 'box'])]
        data_types = ['point']
    else:
        point_inds = random.sample(list(range(label_mask.size(1))), round(label_mask.size(1) / 2))
        data_types = ['point' if idx in point_inds else 'box' for idx in range(label_mask.size(1))]
 
    point_coords, point_labels, point_frames = [], [], []
    prompt_frame_idx = 0
    for obj_idx in range(label_mask.size(1)):
        sample_pool = label_mask.any(dim=(-1, -2))[:, obj_idx].nonzero()[:, 0].tolist()
        area = label_mask[prompt_frame_idx, None, obj_idx, None]
        point_frames.append(prompt_frame_idx)

        if 'point' in data_types[obj_idx]:
            obj_point_coords, obj_point_labels = get_next_point(area, None, 'center', positive_only=True)
            assert (obj_point_labels == 1).all(), obj_point_labels
        elif 'box' in data_types[obj_idx]:
            obj_point_coords, obj_point_labels = sample_box_points(area, noise=-1)
            assert ((obj_point_labels == 2) + (obj_point_labels == 3)).all(), obj_point_labels
        else:
            raise ValueError(f"unknown data type: {data_types[obj_idx]}")

        point_coords.append(obj_point_coords)
        point_labels.append(obj_point_labels)

    point_frames = [torch.LongTensor(point_frames)]

    data['frames'] = [sam2_transform(frames).to(model.sam2.dtype)]
    data['frame_size'] = [frames.shape[1:3]]
    # data['refer_mask'] = [refer_mask]
    data['point_coords'] = [[p.to(device) for p in point_coords]]
    data['point_labels'] = [[p.to(device) for p in point_labels]]
    data['point_frames'] = [[p.to(device) for p in point_frames]]
    
    output_ids = model.generate(
        **data.to(device),
        do_sample=False,
        temperature=None,
        top_k=None,
        top_p=None,
        repetition_penalty=None,
        max_new_tokens=512)

    assert data.input_ids.size(0) == output_ids.size(0) == 1
    output_ids = output_ids[0, data.input_ids.size(1):]

    if output_ids[-1] == processor.tokenizer.eos_token_id:
        output_ids = output_ids[:-1]

    response = processor.decode(output_ids, clean_up_tokenization_spaces=False)
    print(f'\n\033[1;32m   Response:\033[0m {response}')

    response = response.replace(' [0]', '[0]').replace('[0] ', '[0]').replace('[0]', '<REGION>')
    entities = []
    for m in re.finditer(re.escape('<REGION>'), response):
        entities.append(dict(entity='region', start=m.start(), end=m.end(), color="#f85050"))

    answer = dict(text=response, entities=entities)

    print("answer:", answer)

    if len(model.seg) >= 1:
        print("num masks:", len(model.seg))
        # print(model.seg[0].shape) # torch.Size([1, 1, 480, 640])
        imgs = draw_mask(frames, model.seg)

        nncore.mkdir(args.output_dir)

        path = nncore.join(args.output_dir, f"{nncore.pure_name(args.media_path)}.{'gif' if len(imgs) > 1 else 'png'}")
        print(f'\033[1;32mOutput Path:\033[0m {path}')
        iio.imwrite(path, imgs, duration=100, loop=0)