utils/reason_aff_dataset.py

import glob
import json
import os
import random

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from transformers import CLIPImageProcessor

from model.llava import conversation as conversation_lib
from model.segment_anything.utils.transforms import ResizeLongestSide

from .data_processing import get_mask_from_json
from .utils import (ANSWER_LIST, DEFAULT_IMAGE_TOKEN,
                    EXPLANATORY_QUESTION_LIST, LONG_QUESTION_LIST,
                    SHORT_QUESTION_LIST)
from PIL import Image

import pickle


AFFORDANCE_QUESTION_LIST = [
    DEFAULT_IMAGE_TOKEN + "\n" + "Can you segment the affordance map of {class_name} in this image?",
    DEFAULT_IMAGE_TOKEN + "\n" + "Please segment the affordance map of {class_name} in this image.",
    DEFAULT_IMAGE_TOKEN
    + "\n"
    + "What is the affordance map of {class_name} in this image? Please respond with segmentation mask.",
    DEFAULT_IMAGE_TOKEN
    + "\n"
    + "What is the affordance map of {class_name} in this image? Please output segmentation mask.",
]


class ReasonAffDataset(torch.utils.data.Dataset):
    pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
    pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
    img_size = 1024
    ignore_label = 255

    def __init__(
        self,
        base_image_dir,
        tokenizer,
        vision_tower,
        samples_per_epoch=500 * 8 * 2 * 10,
        precision: str = "fp32",
        image_size: int = 224,
        num_classes_per_sample: int = 3,
        exclude_val=False,
        reason_aff_data="handal_hard_reasoning",
        reason_aff_sample_ratio=[1],
        explanatory=0.1,
    ):
        self.exclude_val = exclude_val
        self.reason_aff_data = reason_aff_data
        reason_aff_sample_ratio = np.array(reason_aff_sample_ratio)
        self.reason_aff_sample_ratio = reason_aff_sample_ratio / reason_aff_sample_ratio.sum()
        self.samples_per_epoch = samples_per_epoch
        self.explanatory = explanatory
        self.num_classes_per_sample = num_classes_per_sample

        self.base_image_dir = base_image_dir
        self.image_size = image_size
        self.tokenizer = tokenizer
        self.precision = precision
        self.transform = ResizeLongestSide(image_size)
        self.clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower)

        self.short_question_list = SHORT_QUESTION_LIST
        self.affordance_question_list = AFFORDANCE_QUESTION_LIST
        self.long_question_list = LONG_QUESTION_LIST
        self.answer_list = ANSWER_LIST

        reason_aff_datas = reason_aff_data.split("||")
        self.data2list = {}
        self.object_ids = {}
        for ds in reason_aff_datas:
            if ds == "handal_hard_reasoning" or ds == "egoobjects_easy_reasoning" or ds == "egoobjects_hard_reasoning":
                pkl_path = os.path.join(base_image_dir, f'{ds}_train.pkl')
                images = {}
                labels = {}
                questions = {}
                answers = {}
                with open(pkl_path, 'rb') as f:
                    aff_datas = pickle.load(f)
                for aff_data in aff_datas:
                    if aff_data['task_object_class'] not in images:
                        images[aff_data['task_object_class']] = []
                        labels[aff_data['task_object_class']] = []
                        questions[aff_data['task_object_class']] = []
                        answers[aff_data['task_object_class']] = []
                    images[aff_data['task_object_class']].append(aff_data['frame_path'])
                    labels[aff_data['task_object_class']].append(aff_data['mask_path'])
                    questions[aff_data['task_object_class']].append(aff_data['question'])
                    answers[aff_data['task_object_class']].append(aff_data['answer'])
                # keep same numbers of samples for each class
                for k in images.keys():
                    assert len(images[k]) == len(labels[k])
                self.data2list[ds] = (images, labels, questions, answers)
                print(f"categories of {ds}: ", images.keys())
                print(f"number of {ds} samples: ", len(aff_datas))
            else:
                raise ValueError(f"Unsupported affordance segmentation dataset: {ds}")

    def __len__(self):
        return self.samples_per_epoch

    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
        """Normalize pixel values and pad to a square input."""
        # Normalize colors
        x = (x - self.pixel_mean) / self.pixel_std

        # Pad
        h, w = x.shape[-2:]
        padh = self.img_size - h
        padw = self.img_size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x

    def __getitem__(self, idx):
        ds = np.random.choice(list(self.data2list.keys()), p=self.reason_aff_sample_ratio)

        images, labels, my_questions, my_answers = self.data2list[ds]
        class_name = random.choice(list(images.keys()))
        idx = random.randint(0, len(images[class_name]) - 1)
        image_path = images[class_name][idx]
        label_path = labels[class_name][idx]
        my_question = my_questions[class_name][idx]
        my_answer = my_answers[class_name][idx]

        # load image and prepare input for clip and sam
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        ori_size = image.shape[:2]
        # preprocess image for clip
        image_clip = self.clip_image_processor.preprocess(image, return_tensors="pt")[
            "pixel_values"
        ][0]

        image = self.transform.apply_image(image)  # preprocess image for sam
        resize = image.shape[:2]
        image = self.preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous())

        # load class names
        sampled_classes = [class_name]

        # load label
        label = Image.open(label_path)
        label = np.array(label)
        label = torch.from_numpy(label).long()
        masks = []
        if ds == 'graspnet':
            object_id = self.object_ids[ds][class_name][idx]
            # if data is from graspnet and object_id exists, use the mask of the object_id
            if object_id is None:
                for _ in range(len(sampled_classes)):
                    masks.append(label > 0)
            else:
                for _ in range(len(sampled_classes)):
                    masks.append(label == object_id)
        else:
            for _ in range(len(sampled_classes)):
                masks.append(label > 0)
        masks = torch.stack(masks, dim=0)

        questions = []
        answers = []
        for sampled_cls in sampled_classes:
            text = sampled_cls

            # assert len(text.split("||")) == 1
            # question_template = random.choice(self.affordance_question_list)
            # questions.append(question_template.format(class_name=text.lower()))
            #
            # answers.append(random.choice(self.answer_list))
            questions.append(DEFAULT_IMAGE_TOKEN + "\n" + "You are an embodied robot. " + my_question)
            # answers.append(my_answer + " [SEG].")
            answers.append(my_answer + " [AFF].")

        conversations = []
        conv = conversation_lib.default_conversation.copy()

        i = 0
        while i < len(questions):
            conv.messages = []
            conv.append_message(conv.roles[0], questions[i])
            conv.append_message(conv.roles[1], answers[i])
            conversations.append(conv.get_prompt())
            i += 1

        return (
            image_path,
            image,
            image_clip,
            conversations,
            masks,
            label,
            resize,
            questions,
            sampled_classes,
        )


class ReasonAffValDataset(torch.utils.data.Dataset):
    pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
    pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
    img_size = 1024
    ignore_label = 255

    def __init__(
        self,
        base_image_dir,
        tokenizer,
        vision_tower,
        val_dataset,
        image_size=1024,
    ):
        self.base_image_dir = base_image_dir.replace("/lisa_data", "")
        # splits = val_dataset.split("|")
        # ds, split = splits
        ds = val_dataset

        self.images = []
        self.labels = []
        self.questions = []
        self.answers = []
        self.class_ids = []
        self.class_names = []
        pkl_path = os.path.join(self.base_image_dir, f'{ds}_val.pkl')
        with open(pkl_path, 'rb') as f:
            reason_datas = pickle.load(f)
        for reason_data in reason_datas:
            # one image is broken in 3doi_easy_reasoning_val.pkl, so skip it
            if 'EK_frame_0000040462.jpg' in reason_data['frame_path']:
                continue
            self.images.append(reason_data['frame_path'])
            self.labels.append(reason_data['mask_path'])
            self.questions.append(reason_data['question'])
            self.answers.append(reason_data['answer'])
            self.class_ids.append(None)
            self.class_names.append(reason_data['task_object_class'])

        self.ds = ds
        self.image_size = image_size
        self.tokenizer = tokenizer
        self.transform = ResizeLongestSide(image_size)
        self.clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower)

    def __len__(self):
        return len(self.images)

    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
        """Normalize pixel values and pad to a square input."""
        # Normalize colors
        x = (x - self.pixel_mean) / self.pixel_std

        # Pad
        h, w = x.shape[-2:]
        padh = self.img_size - h
        padw = self.img_size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x

    def __getitem__(self, idx):

        # load image
        image_path = self.images[idx]
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # preprocess image for clip
        image_clip = self.clip_image_processor.preprocess(image, return_tensors="pt")[
            "pixel_values"
        ][0]

        # preprocess image for sam
        image = self.transform.apply_image(image)
        resize = image.shape[:2]
        image = self.preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous())

        # load class names
        sampled_sents = [self.class_names[idx]]

        # load label
        label_path = self.labels[idx]
        label = Image.open(label_path)
        label = np.array(label)
        label = torch.from_numpy(label).long()
        masks = []
        class_id = self.class_ids[idx]
        # if data object_id exists, use the mask of the object_id
        if class_id is None:
            for _ in range(len(sampled_sents)):
                masks.append(label > 0)
        else:
            for _ in range(len(sampled_sents)):
                masks.append(label == class_id)
        masks = torch.stack(masks, dim=0)

        # load question and answer
        my_question = self.questions[idx]
        my_answer = self.answers[idx]

        conversations = []
        conv = conversation_lib.default_conversation.copy()
        i = 0
        while i < len(sampled_sents):
            conv.messages = []
            text = sampled_sents[i].strip()

            conv.append_message(
                conv.roles[0],
                DEFAULT_IMAGE_TOKEN + "\n" + "You are an embodied robot. " + "{}".format(my_question),
            )
            conv.append_message(conv.roles[1], my_answer + " [AFF].")
            conversations.append(conv.get_prompt())
            i += 1

        inference = True

        return (
            image_path,
            image,
            image_clip,
            conversations,
            masks,
            label,
            resize,
            None,
            None,
            inference,
        )