projects/sa2va/datasets/data_utils.py

import copy
from PIL import Image
import cv2
import numpy as np
import torch
from typing import Dict, List, Sequence
from torch.nn.utils.rnn import pad_sequence
from xtuner.dataset.utils import get_bos_eos_token_ids
from xtuner.utils import IGNORE_INDEX, DEFAULT_PAD_TOKEN_INDEX
from xtuner.registry import BUILDER
from mmengine.logging import print_log
import pycocotools.mask as maskUtils
from torch.utils.data import ConcatDataset as TorchConcatDataset


def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height,
                              image_size):
    best_ratio_diff = float('inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio

def dynamic_preprocess(image,
                       min_num=1,
                       max_num=6,
                       image_size=448,
                       use_thumbnail=False):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height

    target_ratios = {(i, j)
                     for n in range(min_num, max_num + 1)
                     for i in range(1, n + 1) for j in range(1, n + 1)
                     if i * j <= max_num and i * j >= min_num}
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio,
                                                    target_ratios, orig_width,
                                                    orig_height, image_size)

    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = ((i % (target_width // image_size)) * image_size,
               (i // (target_width // image_size)) * image_size,
               ((i % (target_width // image_size)) + 1) * image_size,
               ((i // (target_width // image_size)) + 1) * image_size)
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images

def tokenize_conversation(
        example,
        tokenizer,
        max_length,
):
    """We only support the following three scenarios:

    1. Incremental pretraining dataset.
        example['conversation'] = [
                {
                    'input': '',
                    'output': '### Human: Can you write xxx'
                }
            ]

    2. Single-turn conversation dataset.
        example['conversation'] = [
                {
                    'input': 'Give three tips for staying healthy.',
                    'output': '1.Eat a balanced diet xxx'
                }
            ]

    3. Multi-turn conversation dataset.
        example['conversation'] = [
                {
                    'input': 'Give three tips for staying healthy.',
                    'output': '1.Eat a balanced diet xxx'
                },
                {
                    'input': 'Please expand on the second point.',
                    'output': 'Here is an expanded explanation of the xxx'
                }
            ]
    """
    bos_token_id, eos_token_id = get_bos_eos_token_ids(tokenizer)

    input_ids, labels = [], []
    next_needs_bos_token = True
    for single_turn_conversation in example['conversation']:
        input = single_turn_conversation['input']
        input_encode = tokenizer.encode(input, add_special_tokens=False)
        if next_needs_bos_token:
            input_ids += bos_token_id
            labels += [IGNORE_INDEX] * len(bos_token_id)
        input_ids += input_encode
        labels += [IGNORE_INDEX] * len(input_encode)
        output_with_loss = single_turn_conversation.get(
            'output_with_loss', True)
        output = single_turn_conversation['output']
        output_encode = tokenizer.encode(output, add_special_tokens=False)
        input_ids += output_encode
        if output_with_loss:
            labels += copy.deepcopy(output_encode)
        else:
            labels += [IGNORE_INDEX] * len(output_encode)
        if single_turn_conversation.get('need_eos_token', True):
            next_needs_bos_token = True
            input_ids += eos_token_id
            if output_with_loss:
                labels += copy.deepcopy(eos_token_id)
            else:
                labels += [IGNORE_INDEX] * len(eos_token_id)
        else:
            next_needs_bos_token = False
        sep = single_turn_conversation.get('sep', '')
        if sep != '':
            sep_encode = tokenizer.encode(sep, add_special_tokens=False)
            input_ids += sep_encode
            labels += [IGNORE_INDEX] * len(sep_encode)
            

    if len(input_ids) > max_length:
        input_ids = input_ids[:max_length]
        labels = labels[:max_length]
    return {'input_ids': input_ids, 'labels': labels}



def template_map_fn(example, template):
    conversation = example.get("conversation", [])
    for i, single_turn_conversation in enumerate(conversation):
        input = single_turn_conversation.get("input", "")
        if input is None:
            input = ""
        input_text = template.INSTRUCTION.format(input=input, round=i + 1)
        system = single_turn_conversation.get("system", "")
        if system != "" and system is not None:
            system = template.SYSTEM.format(system=system)
            input_text = system + input_text
        single_turn_conversation["input"] = input_text

        if template.get("SUFFIX", None):
            output_text = single_turn_conversation.get("output", "")
            output_text += template.SUFFIX
            single_turn_conversation["output"] = output_text

        single_turn_conversation["need_eos_token"] = not template.get(
            "SUFFIX_AS_EOS", False
        )
        single_turn_conversation["sep"] = template.get("SEP", "")

    return {"conversation": conversation}


def sa2va_collect_fn(
        instances: Sequence[Dict],
        pad_index: int = DEFAULT_PAD_TOKEN_INDEX,
        return_hf_format: bool = False,
        use_varlen_attn: bool = False
):
    assert not return_hf_format, "return_hf_format is not supported yet."
    assert not use_varlen_attn, "use_varlen_attn is not supported yet."

    input_ids, labels = [], []

    has_image = any(inst.get('pixel_values') is not None for inst in instances)
    has_pe = any(inst.get('image_grid_thw', None) is not None for inst in instances)
    has_grounding_image = any(inst.get('g_pixel_values') is not None for inst in instances)
    has_mask = any(inst.get('masks') is not None for inst in instances)


    has_vp = any(inst.get('vp_overall_mask') is not None for inst in instances)
    has_prompt_mask = any(inst.get('prompt_masks') is not None for inst in instances)
    assert has_vp and has_prompt_mask or not has_vp and not has_prompt_mask, \
    f"Inconsistent presence of visual prompts and prompt masks {has_vp} {has_prompt_mask}"

    pixel_values = []
    frames_per_batch = []
    
    image_grid_thw = []
    grounding_pixel_values = []
    object_masks = []
    vp_overall_mask = []
    prompt_masks = []
    for example in instances:
        input_ids.append(torch.LongTensor(example['input_ids']))
        labels.append(torch.LongTensor(example['labels']))

        if has_image:
            pixel_values.append(example['pixel_values'])
            if has_pe:
                image_grid_thw.append(example['image_grid_thw'])
            if has_vp:
                if 'vp_overall_mask' in example.keys() and example['vp_overall_mask'] is not None:
                    vp_overall_mask.append(example['vp_overall_mask'])
                else:
                    vp_overall_mask.append(torch.Tensor([False] * len(example['pixel_values'])))
        
        if has_grounding_image and 'g_pixel_values' in example.keys():
            if isinstance(example['g_pixel_values'], list):
                grounding_pixel_values += example['g_pixel_values']
                frames_per_batch.append(len(example['g_pixel_values']))
            else:
                grounding_pixel_values.append(example['g_pixel_values'])
                frames_per_batch.append(1)

        if has_mask:
            if 'masks' in example.keys() and example['masks'] is not None:
                if isinstance(example['masks'], list):
                    if isinstance(example['masks'][0], np.ndarray):
                        _masks = np.stack(example['masks'], axis=0)
                        _masks = torch.from_numpy(_masks)
                        object_masks.append(_masks)
                    else:
                        object_masks.append(torch.stack(example['masks'], dim=0))
                else:
                    object_masks.append(example['masks'])

        if has_prompt_mask:
            if 'prompt_masks' in example.keys():
                prompt_masks.append(example['prompt_masks'])

    ori_length = [len(ids) for ids in input_ids]
    if len(instances) > 1:
        input_ids = pad_sequence(
            input_ids, batch_first=True, padding_value=pad_index)
        labels = pad_sequence(
            labels, batch_first=True, padding_value=IGNORE_INDEX)
    else:
        input_ids = torch.stack(input_ids)
        labels = torch.stack(labels)

    attention_mask = torch.zeros_like(input_ids).bool()
    for i, length in enumerate(ori_length):
        attention_mask[i, :length] = True

    bs, seq_len = input_ids.shape
    position_ids = torch.arange(seq_len).unsqueeze(0).long().repeat(bs, 1)

    data_dict = {
        'input_ids': input_ids,
        'attention_mask': attention_mask,
        'position_ids': position_ids,
        'labels': labels
    }

    if has_image:
        data_dict['frames_per_batch'] = frames_per_batch
        data_dict['pixel_values'] = pixel_values
        for pixel_values_per_sample in pixel_values:
            assert isinstance(pixel_values_per_sample, torch.Tensor)
        
        if has_pe:
            data_dict['image_grid_thw'] = image_grid_thw

    if has_vp:
        data_dict['vp_overall_mask'] = torch.cat(vp_overall_mask, dim=0)

    if has_prompt_mask:
        data_dict['prompt_masks'] = prompt_masks

    if has_grounding_image:
        data_dict['g_pixel_values'] = grounding_pixel_values

    if has_mask:
        data_dict['masks'] = object_masks

    return {'data': data_dict, 'data_samples': None}


def sa2va_collect_fn_multitask(
        instances: Sequence[Dict],
        pad_index: int = DEFAULT_PAD_TOKEN_INDEX,
        return_hf_format: bool = False,
        use_varlen_attn: bool = False
):
    assert not return_hf_format, "return_hf_format is not supported yet."
    assert not use_varlen_attn, "use_varlen_attn is not supported yet."

    g_pixel_values = []
    masks = []
    src = []
    meta = []
    images_star = []
    images_without_star = []
    for ex in instances:
        if "g_pixel_values" not in ex or "masks" not in ex:
            raise ValueError("Expected g_pixel_values and masks in multitask example")
        g_pixel_values.append(ex["g_pixel_values"])
        masks.append(ex["masks"])
        src.append(ex.get("src", "unknown"))
        meta.append(ex.get("meta", None))
        if ex.get("images_star", None) is not None:
            images_star.append(ex["images_star"])
        if ex.get("images_without_star", None) is not None:
            images_without_star.append(ex["images_without_star"])

    tasks_out: Dict[str, Dict[str, torch.Tensor]] = {}
    for task_name in ["star", "referring", "vqa"]:
        task_items = []
        base_indices = []
        convs = []
        questions = []
        for base_i, ex in enumerate(instances):
            t = ex.get("tasks", {}).get(task_name, None)
            if t is None:
                continue
            task_items.append(t)
            base_indices.append(base_i)
            if "convs" in t and t["convs"] is not None:
                convs.append(t["convs"])
            else:
                convs.append("")
            if "question" in t and t["question"] is not None:
                questions.append(str(t["question"]))
            else:
                questions.append("")

        if len(task_items) == 0:
            continue

        input_ids = [torch.LongTensor(t["input_ids"]) for t in task_items]
        labels = [torch.LongTensor(t["labels"]) for t in task_items]
        pixel_values = [t["pixel_values"] for t in task_items]
        image_grid_thw = [t["image_grid_thw"] for t in task_items]

        ori_length = [len(ids) for ids in input_ids]
        if len(input_ids) > 1:
            input_ids = pad_sequence(input_ids, batch_first=True, padding_value=pad_index)
            labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
        else:
            input_ids = torch.stack(input_ids)
            labels = torch.stack(labels)

        attention_mask = torch.zeros_like(input_ids).bool()
        for i, length in enumerate(ori_length):
            attention_mask[i, :length] = True

        position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0).long().repeat(input_ids.shape[0], 1)

        tasks_out[task_name] = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "labels": labels,
            "pixel_values": pixel_values,
            "image_grid_thw": image_grid_thw,
            "base_indices": torch.LongTensor(base_indices),
            "convs": convs,
            "questions": questions,
        }

    if len(tasks_out) == 0:
        raise ValueError("No tasks found in multitask batch")

    data_dict = {
        "tasks": tasks_out,
        "g_pixel_values": g_pixel_values,
        "masks": masks,
        "frames_per_batch": [1 for _ in range(len(instances))],
        "src": src,
        "meta": meta,
    }
    if len(images_star) > 0:
        data_dict["images_star"] = torch.stack(images_star, dim=0)
    else:
        data_dict["images_star"] = None
    if len(images_without_star) > 0:
        data_dict["images_without_star"] = torch.stack(images_without_star, dim=0)
    else:
        data_dict["images_without_star"] = None
    return {"data": data_dict, "data_samples": None}


def sam2_path_patch(video_path, anno_path):
    if 'sav_train' in video_path:
        path_parts = video_path.split('/')
        sav_train_idx = None
        duplicate_idx = None
        for i, part in enumerate(path_parts):
            if part == 'sav_train':
                assert sav_train_idx is None, "Multiple 'sav_train' directories found."
                sav_train_idx = i
        if sav_train_idx is not None:
            if path_parts[sav_train_idx - 1] == path_parts[sav_train_idx + 1]:
                duplicate_idx = sav_train_idx - 1
        
        if duplicate_idx is not None:
            del path_parts[duplicate_idx]
            video_path = '/'.join(path_parts)

            anno_parts = anno_path.split('/')
            del anno_parts[duplicate_idx]
            anno_path = '/'.join(anno_parts)
    return video_path, anno_path


def get_video_frames(video_path) -> List[np.ndarray]:
    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        print("Error: Cannot open video file.")
        return []

    frames = []

    frame_id = 0
    while True:
        ret, frame = cap.read()

        if not ret:
            break

        frames.append(frame)

        frame_id += 1

    cap.release()
    return frames

def decode_masklet(masklet):
    masks = []
    for _rle in masklet:
        mask = maskUtils.decode(_rle)
        masks.append(mask)
    return masks


def opencvimg_to_pil(image: np.ndarray) -> Image.Image:
    """Convert an OpenCV image (BGR) to a PIL image (RGB)."""
    image = image[:, :, ::-1]  # Convert BGR to RGB
    return Image.fromarray(image).convert('RGB')


class ConcatDatasetSa2VA(TorchConcatDataset):

    def __init__(self, datasets:List[dict]):
        datasets_instance = []
        for cfg in datasets:
            datasets_instance.append(BUILDER.build(cfg))
        super().__init__(datasets=datasets_instance)

        print_log(
            f'Initialized ConcatDataset with {len(datasets)} datasets.'
        )
        for dataset in self.datasets:
            print_log(f'{repr(dataset.name)}')
            print_log(f'------Number of samples: {len(dataset)}')
            print_log(f'------Real Length: {dataset.real_len()}')

    def __repr__(self):
        main_str = 'Dataset as a concatenation of multiple datasets. \n'
        main_str += ',\n'.join(
            [f'{repr(dataset)}' for dataset in self.datasets])
        return main_str