model/dataset.py

import os
from typing import List, Union
import json
import cv2
import lmdb
import numpy as np
import pyarrow as pa
import torch
from torch.utils.data import Dataset
from copy import deepcopy
import functools
from skimage.measure import regionprops
from shapely.geometry import Polygon
from skimage.draw import polygon
from skimage.filters import gaussian
import matplotlib.pyplot as plt


from .simple_tokenizer import SimpleTokenizer as _Tokenizer
from .OCID_sub_class_dict import cnames, colors, subnames, sub_to_class
from .augmentation import DataAugmentor

info = {
    'refcoco': {
        'train': 42404,
        'val': 3811,
        'val-test': 3811,
        'testA': 1975,
        'testB': 1810
    },
    'refcoco+': {
        'train': 42278,
        'val': 3805,
        'val-test': 3805,
        'testA': 1975,
        'testB': 1798
    },
    'refcocog_u': {
        'train': 42226,
        'val': 2573,
        'val-test': 2573,
        'test': 5023
    },
    'refcocog_g': {
        'train': 44822,
        'val': 5000,
        'val-test': 5000
    },
    'cocostuff': {
        "train": 965042,
        'val': 42095,
        'val-test': 42095
    }
}
_tokenizer = _Tokenizer()


def tokenize(texts: Union[str, List[str]],
             context_length: int = 77,
             truncate: bool = False) -> torch.LongTensor:
    """
    Returns the tokenized representation of given input string(s)

    Parameters
    ----------
    texts : Union[str, List[str]]
        An input string or a list of input strings to tokenize

    context_length : int
        The context length to use; all CLIP models use 77 as the context length

    truncate: bool
        Whether to truncate the text in case its encoding is longer than the context length

    Returns
    -------
    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
    """
    if isinstance(texts, str):
        texts = [texts]

    sot_token = _tokenizer.encoder["<|startoftext|>"]
    eot_token = _tokenizer.encoder["<|endoftext|>"]
    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token]
                  for text in texts]
    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)

    for i, tokens in enumerate(all_tokens):
        if len(tokens) > context_length:
            if truncate:
                tokens = tokens[:context_length]
                tokens[-1] = eot_token
            else:
                raise RuntimeError(
                    f"Input {texts[i]} is too long for context length {context_length}"
                )
        result[i, :len(tokens)] = torch.tensor(tokens)

    return result


def loads_pyarrow(buf):
    """
    Args:
        buf: the output of `dumps`.
    """
    return pa.deserialize(buf)



class RefOCIDGraspDataset(Dataset):
    def __init__(self, root_path, input_size, word_length, mode="train"):
        super().__init__()
        json_path = os.path.join(root_path, f"{mode}_expressions.json")
        with open(json_path, "r") as f:
            self.meta_data = json.load(f)
        
        self.root_path = root_path
        self.keys = list(self.meta_data.keys())
        self.input_size = (input_size, input_size)
        self.word_length = word_length
        self.mode = mode
        
        self.cls_names = cls_names
        self.colors = colors
        self.mean = torch.tensor([0.48145466, 0.4578275,
                                  0.40821073]).reshape(3, 1, 1)
        self.std = torch.tensor([0.26862954, 0.26130258,
                                 0.27577711]).reshape(3, 1, 1)
        
        self.target_save_dir = os.path.join("./", mode)
        os.makedirs(self.target_save_dir, exist_ok=True)

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

    def visualization(self, rgb, depth, masks, rects, grasp_masks, bbox, obj_cls, sentence):
        cls_list = list(self.cls_names.keys())

        print(rgb.shape, depth.shape, len(rects), masks.shape)
        mask_color = self.colors[str(obj_cls)]
        rgb_with_grasp = deepcopy(rgb)

        color_masks = np.repeat(masks[:, :, np.newaxis], 3, axis=-1) * mask_color
        rgb = (rgb * 255).astype(np.uint8)
        img_fuse = (rgb * 0.3 + color_masks * 0.7).astype(np.uint8)

        x1, y1, x2, y2 = bbox
        cv2.rectangle(img_fuse, (x1, y1), (x2, y2), (255, 0, 0), 2)


        print(f'\nimg shape: {rgb.shape}')
        print('----------------boxes----------------')
        for rect in rects:
            print(rect)
        print('----------------labels---------------')
        print([cls_list[int(i)] for i in [obj_cls]], '\n')

        for rect in rects:
            name = cls_list[int(obj_cls)]
            color = self.colors[str(obj_cls)]
            center_x, center_y, width, height, theta, cls_id = rect
            box = ((center_x, center_y), (width, height), -(theta+180))
            box = cv2.boxPoints(box)
            box = np.int0(box)
            cv2.drawContours(rgb_with_grasp, [box], 0, color.tolist(), 2)
        

        fig = plt.figure(figsize=(25, 10))

        ax = fig.add_subplot(2, 4, 1)
        ax.imshow(rgb)
        ax.set_title('RGB')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 2)
        ax.imshow(depth, cmap='gray')
        ax.set_title('Depth')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 3)
        ax.imshow(img_fuse)
        ax.set_title('Masks & Bboxes')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 4)
        ax.imshow(rgb_with_grasp)
        ax.set_title('Grasps')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 5)
        plot = ax.imshow(grasp_masks[0], cmap='jet', vmin=0, vmax=1)
        ax.set_title('Quality')
        ax.axis('off')
        plt.colorbar(plot)

        ax = fig.add_subplot(2, 4, 6)
        plot = ax.imshow(grasp_masks[1], cmap='jet', vmin=0, vmax=1)
        ax.set_title('Quality')
        ax.axis('off')
        plt.colorbar(plot)

        ax = fig.add_subplot(2, 4, 7)
        plot = ax.imshow(grasp_masks[2], cmap='rainbow', vmin=-np.pi / 2, vmax=np.pi / 2)
        ax.set_title('Angle')
        ax.axis('off')
        plt.colorbar(plot)

        ax = fig.add_subplot(2, 4, 8)
        plot = ax.imshow(grasp_masks[3], cmap='jet', vmin=0, vmax=1)
        ax.set_title('Width')
        ax.axis('off')
        plt.colorbar(plot)

        plt.suptitle(f"{sentence}", fontsize=20)
        plt.tight_layout()
        plt.savefig("./visualization.png")

    # @functools.lru_cache(maxsize=None)
    def _load_bbox(self, bbox):
        bbox = bbox.replace("[", "").replace("]", "").split(",")
        x, y, w, h = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])

        return [x, y, w, h]

    # @functools.lru_cache(maxsize=None)
    def _load_rgb(self, path):
        image = cv2.imread(os.path.join(self.root_path, path))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
        return image
    
    # @functools.lru_cache(maxsize=None)
    def _load_depth(self, path, factor=1000.):
        depth = cv2.imread(os.path.join(self.root_path, path), cv2.IMREAD_UNCHANGED) / factor
        depth = 1 - (depth / np.max(depth))

        return depth
    
    # @functools.lru_cache(maxsize=None)
    def _load_annos(self, path, target_cls):
        cls_annos_path = os.path.join(os.path.join(self.root_path, path), str(target_cls))
        file_id = cls_annos_path.split("/")[-2]
        cls_annos_path = os.path.join(cls_annos_path, f"{file_id}.txt")

        grasps_list = []
        with open(cls_annos_path, 'r') as f:
            points_list = []
            for count, line in enumerate(f):
                line = line.rstrip()
                [x, y] = line.split(' ')

                x = float(x)
                y = float(y)

                pt = (x, y)
                points_list.append(pt)

                if len(points_list) == 4:
                    p1, p2, p3, p4 = points_list
                    center_x = (p1[0] + p3[0]) / 2
                    center_y = (p1[1] + p3[1]) / 2
                    width  = np.sqrt((p1[0] - p4[0]) * (p1[0] - p4[0]) + (p1[1] - p4[1]) * (p1[1] - p4[1]))
                    height = np.sqrt((p1[0] - p2[0]) * (p1[0] - p2[0]) + (p1[1] - p2[1]) * (p1[1] - p2[1]))
                    
                    # @NOTE
                    # Along x+ is 0 degree, increase by rotating anti-clockwise
                    # If you want to use opencv boxPoints & drawContours to visualize grasps
                    # Remember to take -theta as param :-)
                    theta = np.arctan2(p4[0] - p1[0], p4[1] - p1[1]) * 180 / np.pi
                    if theta > 0:
                        theta = theta-90
                    else:
                        theta = theta+90


                    grasps_list.append([center_x, center_y, width, height, theta, int(target_cls)])
                    points_list = []

        return grasps_list

    # @functools.lru_cache(maxsize=None)
    def _load_mask(self, path, target_cls):
        masks = cv2.imread(os.path.join(self.root_path, path), cv2.IMREAD_UNCHANGED)
        if target_cls == -1:
            # for loading instance masks
            return masks
        else:
            # For loading target semantic masks
            target_masks = (masks == target_cls)
            return target_masks

    def _match_masks_with_ref(self, bbox, ins_masks, masks):
        # Preparing bounding box for calculating IoU
        x1, y1, x2, y2 = bbox
        w, h = (x2-x1), (y2-y1)
        vertices = [[x1, y1], [x1+w, y1], [x2, y2], [x1, y1+h]]
        poly1 = Polygon(vertices)

        # Keep only the instance masks with correct class label
        ins_masks = ins_masks * masks
        ins_props = regionprops(ins_masks)

        max_iou = 0.0
        ins_idx = 0

        for ins in ins_props:
            _x1, _y1, _x2, _y2 = ins.bbox[1], ins.bbox[0], ins.bbox[3], ins.bbox[2]
            _w, _h = (_x2-_x1), (_y2-_y1)
            ins_vertices = [[_x1, _y1], [_x1+_w, _y1], [_x2, _y2], [_x1, _y1+_h]]
            poly2 = Polygon(ins_vertices)

            if poly1.intersects(poly2): 
                intersect = poly1.intersection(poly2).area
                union = poly1.union(poly2).area
                iou = intersect/union

                if iou > max_iou:
                    max_iou = iou
                    ins_idx = ins.label
        
        ins_masks = (ins_masks == ins_idx)

        return ins_masks

    def _match_grasps_with_ref(self, rects, ins_masks):
        # Check if the center of grasp falls in the target instance mask
        grasps = []
        for rect in rects:
            c_x, c_y = int(rect[0]), int(rect[1])
            if ins_masks[c_y, c_x]:
                grasps.append(rect)
        
        return grasps

    def _filter_grasps(self, rects):
        angles = []
        for rect in rects:
            angles.append(rect[4])
            
    
    def _generate_grasp_masks(self, grasps, width, height):
        pos_out = np.zeros((height, width))
        ang_out = np.zeros((height, width))
        wid_out = np.zeros((height, width))
        for rect in grasps:
            center_x, center_y, w_rect, h_rect, theta, cls_id = rect
            width_factor = float(100)

            # Get 4 corners of rotated rect
            # Convert from our angle represent to opencv's
            r_rect = ((center_x, center_y), (w_rect/2, h_rect), -(theta+180))
            box = cv2.boxPoints(r_rect)
            box = np.int0(box)

            rr, cc = polygon(box[:, 0], box[:,1])

            mask_rr = rr < width
            rr = rr[mask_rr]
            cc = cc[mask_rr]

            mask_cc = cc < height
            cc = cc[mask_cc]
            rr = rr[mask_cc]


            pos_out[cc, rr] = 1.0
            if theta < 0:
                ang_out[cc, rr] = int(theta + 180)
            else:
                ang_out[cc, rr] = int(theta)
            # Adopt width normalize accoding to class 
            wid_out[cc, rr] = np.clip(w_rect, 0.0, width_factor) / width_factor
        
        qua_out = (gaussian(pos_out, 3, preserve_range=True) * 255).astype(np.uint8)
        pos_out = (pos_out * 255).astype(np.uint8)
        ang_out = ang_out.astype(np.uint8)
        wid_out = (gaussian(wid_out, 3, preserve_range=True) * 255).astype(np.uint8)
        
        return [pos_out, qua_out, ang_out, wid_out]
    
    def getTransformMat(self, img_size, inverse=False):
        ori_h, ori_w = img_size
        inp_h, inp_w = self.input_size
        scale = min(inp_h / ori_h, inp_w / ori_w)
        new_h, new_w = ori_h * scale, ori_w * scale
        bias_x, bias_y = (inp_w - new_w) / 2., (inp_h - new_h) / 2.

        src = np.array([[0, 0], [ori_w, 0], [0, ori_h]], np.float32)
        dst = np.array([[bias_x, bias_y], [new_w + bias_x, bias_y],
                        [bias_x, new_h + bias_y]], np.float32)

        mat = cv2.getAffineTransform(src, dst)
        if inverse:
            mat_inv = cv2.getAffineTransform(dst, src)
            return mat, mat_inv
        return mat, None


    def getTransformMat(self, img_size, inverse=False):
        ori_h, ori_w = img_size
        inp_h, inp_w = self.input_size
        scale = min(inp_h / ori_h, inp_w / ori_w)
        new_h, new_w = ori_h * scale, ori_w * scale
        bias_x, bias_y = (inp_w - new_w) / 2., (inp_h - new_h) / 2.

        src = np.array([[0, 0], [ori_w, 0], [0, ori_h]], np.float32)
        dst = np.array([[bias_x, bias_y], [new_w + bias_x, bias_y],
                        [bias_x, new_h + bias_y]], np.float32)

        mat = cv2.getAffineTransform(src, dst)
        if inverse:
            mat_inv = cv2.getAffineTransform(dst, src)
            return mat, mat_inv
        return mat, None


    def convert(self, img, mask=None, grasp_quality_mask=None, grasp_sin_masks=None, grasp_cos_masks=None, grasp_width_mask=None):
        # Image ToTensor & Normalize
        img = torch.from_numpy(img.transpose((2, 0, 1)))
        if not isinstance(img, torch.FloatTensor):
            img = img.float()
        img.div_(255.).sub_(self.mean).div_(self.std)
        # Mask ToTensor
        if mask is not None:
            mask = torch.from_numpy(mask)
            if not isinstance(mask, torch.FloatTensor):
                mask = mask.float()
        
        if grasp_quality_mask is not None:
            grasp_quality_mask = torch.from_numpy(grasp_quality_mask)
            if not isinstance(grasp_quality_mask, torch.FloatTensor):
                grasp_quality_mask = grasp_quality_mask.float()

        if grasp_sin_masks is not None:
            grasp_sin_masks = torch.from_numpy(grasp_sin_masks)
            if not isinstance(grasp_sin_masks, torch.FloatTensor):
                grasp_sin_masks = grasp_sin_masks.float()
        
        if grasp_cos_masks is not None:
            grasp_cos_masks = torch.from_numpy(grasp_cos_masks)
            if not isinstance(grasp_cos_masks, torch.FloatTensor):
                grasp_cos_masks = grasp_cos_masks.float()

        if grasp_width_mask is not None:
            grasp_width_mask = torch.from_numpy(grasp_width_mask)
            if not isinstance(grasp_width_mask, torch.FloatTensor):
                grasp_width_mask = grasp_width_mask.float()
        return img, mask, grasp_quality_mask, grasp_sin_masks, grasp_cos_masks, grasp_width_mask


    def __getitem__(self, index):
        data_dict = {}

        key = self.keys[index]
        ref_data = self.meta_data[key]
        obj_cls = int(self.cls_names[ref_data["class"]])

        # Get path to other data
        scene_path = ref_data["scene_path"]
        depth_path = scene_path.replace("rgb", "depth")
        annos_path = scene_path.replace("rgb", "Annotations_per_class")[:-4]
        masks_path = scene_path.replace("rgb", "seg_mask_labeled_combi")
        ins_masks_path = scene_path.replace("rgb", "seg_mask_instances_combi")

        # Read data
        rgb = self._load_rgb(scene_path) # [0 - 255]
        depth = self._load_depth(depth_path)
        annos = self._load_annos(annos_path, obj_cls)
        masks = self._load_mask(masks_path, obj_cls)
        ins_masks = self._load_mask(ins_masks_path, -1)
        bbox = self._load_bbox(ref_data["bbox"])
        sentence = ref_data["sentence"]

        # ins_masks = ins_masks * masks

        ins_masks = (self._match_masks_with_ref(bbox, ins_masks, masks) * 255).astype(np.uint8)

        grasps = self._match_grasps_with_ref(annos, ins_masks)

        assert rgb.shape[:2] == depth.shape

        height, width = depth.shape

        grasp_masks = self._generate_grasp_masks(grasps, width, height)

        grasp_quality_masks = grasp_masks[1]
        grasp_angle_masks = grasp_masks[2]
        grasp_width_masks = grasp_masks[3]

        # Image transforms
        img_size = rgb.shape[:2]
        mat, mat_inv = self.getTransformMat(img_size, True)
        rgb = cv2.warpAffine(rgb,
                                mat,
                                self.input_size,
                                flags=cv2.INTER_CUBIC,
                                borderValue=[0.48145466 * 255, 0.4578275 * 255, 0.40821073 * 255])
        if self.mode == 'train':
            ins_masks = cv2.warpAffine(ins_masks,
                                    mat,
                                    self.input_size,
                                    flags=cv2.INTER_LINEAR,
                                    borderValue=0.)
            
            grasp_quality_masks = cv2.warpAffine(grasp_quality_masks,
                                    mat,
                                    self.input_size,
                                    flags=cv2.INTER_LINEAR,
                                    borderValue=0.)
            
            grasp_angle_masks = cv2.warpAffine(grasp_angle_masks,
                                    mat,
                                    self.input_size,
                                    flags=cv2.INTER_LINEAR,
                                    borderValue=0.)
            
            grasp_width_masks = cv2.warpAffine(grasp_width_masks,
                                    mat,
                                    self.input_size,
                                    flags=cv2.INTER_LINEAR,
                                    borderValue=0.)
            
            # Normalize and pre-process target masks
            ins_masks = ins_masks / 255.
            grasp_quality_masks = grasp_quality_masks / 255.
            grasp_angle_masks = grasp_angle_masks * np.pi / 180.
            grasp_width_masks = grasp_width_masks / 255.
            grasp_sin_masks = np.sin(2 * grasp_angle_masks)
            grasp_cos_masks = np.cos(2 * grasp_angle_masks)
            
            word_vec = tokenize(sentence, self.word_length, True).squeeze(0)

            # self.visualization(rgb, depth, ins_masks, grasps, (grasp_quality_masks, grasp_sin_masks, grasp_cos_masks, grasp_width_masks), bbox, obj_cls, sentence)

            rgb, ins_masks, grasp_quality_masks, grasp_sin_masks, grasp_cos_masks, grasp_width_masks = self.convert(
                rgb, ins_masks, grasp_quality_masks, grasp_sin_masks, grasp_cos_masks, grasp_width_masks
            )

            # np.savez(
            #     os.path.join(self.target_save_dir, f"{key}.npz"),
            #     img=rgb.cpu().numpy(),
            #     word_vec=word_vec.cpu().numpy(),
            #     ins_masks=ins_masks.cpu().numpy(),
            #     grasp_quality_masks=grasp_quality_masks.cpu().numpy(),
            #     grasp_sin_masks=grasp_sin_masks.cpu().numpy(),
            #     grasp_cos_masks=grasp_cos_masks.cpu().numpy(),
            #     grasp_width_masks=grasp_width_masks.cpu().numpy(),
            #     grasps=grasps
            # )


            return rgb, (ins_masks, grasp_quality_masks, grasp_sin_masks, grasp_cos_masks, grasp_width_masks), word_vec
        
        elif self.mode == 'val':
            word_vec = tokenize(sentence, self.word_length, True).squeeze(0)

            ins_masks = ins_masks / 255.
            grasp_quality_masks = grasp_quality_masks / 255.
            grasp_angle_masks = grasp_angle_masks * np.pi / 180.
            grasp_width_masks = grasp_width_masks / 255.
            grasp_sin_masks = np.sin(2 * grasp_angle_masks)
            grasp_cos_masks = np.cos(2 * grasp_angle_masks)

            rgb = self.convert(rgb)[0]
            params = {
                # 'mask_dir': mask_dir,
                'inverse': mat_inv,
                'ori_size': np.array(img_size)
            }
            # np.savez(
            #     os.path.join(self.target_save_dir, f"{key}.npz"),
            #     img=rgb.cpu().numpy(),
            #     word_vec=word_vec.cpu().numpy(),
            #     ins_masks=ins_masks,
            #     grasp_quality_masks=grasp_quality_masks,
            #     grasp_sin_masks=grasp_sin_masks,
            #     grasp_cos_masks=grasp_cos_masks,
            #     grasp_width_masks=grasp_width_masks,
            #     grasps=grasps,
            #     params=params
            # )
            return rgb, (ins_masks, grasp_quality_masks, grasp_sin_masks, grasp_cos_masks, grasp_width_masks), word_vec, params
        elif self.mode == 'test':
            rgb = self.convert(rgb)[0]
            word_vec = tokenize(sentence, self.word_length, True).squeeze(0)
            params = {
                # 'ori_img': ori_img,
                # 'seg_id': seg_id,
                # 'mask_dir': mask_dir,
                'inverse': mat_inv,
                'ori_size': np.array(img_size),
                'sents': sentence

            }
            # np.savez(
            #     os.path.join(self.target_save_dir, f"{key}.npz"),
            #     img=rgb.cpu().numpy(),
            #     word_vec=word_vec.cpu().numpy(),
            #     params=params
            # )
            return rgb, word_vec, params



class GraspTransforms:
    # Class for converting cv2-like rectangle formats and generate grasp-quality-angle-width masks

    def __init__(self, width_factor=100, width=640, height=480):
        self.width_factor = width_factor
        self.width = width 
        self.height = height

    def __call__(self, grasp_rectangles, target):
        # grasp_rectangles: (M, 4, 2)
        M = grasp_rectangles.shape[0]
        p1, p2, p3, p4 = np.split(grasp_rectangles, 4, axis=1)
        
        center_x = (p1[..., 0] + p3[..., 0]) / 2
        center_y = (p1[..., 1] + p3[..., 1]) / 2
        
        width  = np.sqrt((p1[..., 0] - p4[..., 0]) * (p1[..., 0] - p4[..., 0]) + (p1[..., 1] - p4[..., 1]) * (p1[..., 1] - p4[..., 1]))
        height = np.sqrt((p1[..., 0] - p2[..., 0]) * (p1[..., 0] - p2[..., 0]) + (p1[..., 1] - p2[..., 1]) * (p1[..., 1] - p2[..., 1]))
        
        theta = np.arctan2(p4[..., 0] - p1[..., 0], p4[..., 1] - p1[..., 1]) * 180 / np.pi
        theta = np.where(theta > 0, theta - 90, theta + 90)

        target = np.tile(np.array([[target]]), (M,1))

        return np.concatenate([center_x, center_y, width, height, theta, target], axis=1)

    def inverse(self, grasp_rectangles):
        boxes = []
        for rect in grasp_rectangles:
            center_x, center_y, width, height, theta = rect[:5]
            box = ((center_x, center_y), (width, height), -(theta+180))
            box = cv2.boxPoints(box)
            box = np.intp(box)
            boxes.append(box)
        return boxes

    def generate_masks(self, grasp_rectangles):
        pos_out = np.zeros((self.height, self.width))
        ang_out = np.zeros((self.height, self.width))
        wid_out = np.zeros((self.height, self.width))
        for rect in grasp_rectangles:
            center_x, center_y, w_rect, h_rect, theta = rect[:5]
            
            # Get 4 corners of rotated rect
            # Convert from our angle represent to opencv's
            r_rect = ((center_x, center_y), (w_rect/2, h_rect), -(theta+180))
            box = cv2.boxPoints(r_rect)
            box = np.intp(box)

            rr, cc = polygon(box[:, 0], box[:,1])

            mask_rr = rr < self.width
            rr = rr[mask_rr]
            cc = cc[mask_rr]

            mask_cc = cc < self.height
            cc = cc[mask_cc]
            rr = rr[mask_cc]
            pos_out[cc, rr] = 1.0
            if theta < 0:
                ang_out[cc, rr] = int(theta + 180)
            else:
                ang_out[cc, rr] = int(theta)
            # Adopt width normalize accoding to class 
            wid_out[cc, rr] = np.clip(w_rect, 0.0, self.width_factor) / self.width_factor
        
        qua_out = (gaussian(pos_out, 3, preserve_range=True) * 255).astype(np.uint8)
        pos_out = (pos_out * 255).astype(np.uint8)
        ang_out = ang_out.astype(np.uint8)
        wid_out = (gaussian(wid_out, 3, preserve_range=True) * 255).astype(np.uint8)
        
        
        return {'pos': pos_out, 
                'qua': qua_out, 
                'ang': ang_out, 
                'wid': wid_out}



class OCIDVLGDataset(Dataset):
    
    """ OCID-Vision-Language-Grasping dataset with referring expressions and grasps """

    def __init__(self, 
                 root_dir,
                 split, 
                 transform_img = None,
                 transform_grasp = GraspTransforms(),
                 input_size = 416,
                 word_length = 20,
                 with_depth = True, 
                 with_segm_mask = True,
                 with_grasp_masks = True,
                 version="multiple"
    ):
        super(OCIDVLGDataset, self).__init__()
        self.root_dir = root_dir
        self.split_dir = os.path.join(root_dir, "data_split")
        self.split_map = {'train': 'train_expressions.json', 
                          'val': 'val_expressions.json',
                          'test': 'test_expressions.json'
                         }
        self.split = split
        self.refer_dir = os.path.join(root_dir, "refer", version)
        
        self.transform_img = transform_img
        self.transform_grasp = transform_grasp
        self.with_depth = with_depth
        self.with_segm_mask = with_segm_mask
        self.with_grasp_masks = with_grasp_masks
        # assert (self.transform_grasp and self.with_grasp_masks) or (not self.transform_grasp and not self.with_grasp_masks)

        self.input_size = (input_size, input_size)
        self.word_length = word_length
        self.mean = torch.tensor([0.48145466, 0.4578275,
                                  0.40821073]).reshape(3, 1, 1)
        self.std = torch.tensor([0.26862954, 0.26130258,
                                 0.27577711]).reshape(3, 1, 1)

        self._load_dicts()
        self._load_split()

    def _load_dicts(self):
        cwd = os.getcwd()
        os.chdir(self.root_dir)
        from .OCID_sub_class_dict import cnames, colors, subnames, sub_to_class
        cnames_inv = {int(v):k for k,v in cnames.items()}
        subnames_inv = {v:k for k,v in subnames.items()}
        self.class_names = cnames 
        self.idx_to_class = cnames_inv
        self.class_instance_names = subnames
        self.idx_to_class_instance = subnames_inv
        self.instance_idx_to_class_idx = sub_to_class
        os.chdir(cwd)

    def _load_split(self):
        refer_data = json.load(open(os.path.join(self.refer_dir, self.split_map[self.split])))
        self.seq_paths, self.img_names, self.scene_ids = [], [], []
        self.bboxes, self.grasps = [], []
        self.sent_to_index, self.sent_indices = {}, []
        self.rgb_paths, self.depth_paths, self.mask_paths = [], [], []
        self.targets, self.sentences, self.semantics, self.objIDs = [], [], [], []
        n = 0
        for item in refer_data['data']:
            seq_path, im_name = item['image_filename'].split(',')
            self.seq_paths.append(seq_path)
            self.img_names.append(im_name)
            self.scene_ids.append(item['image_filename'])
            self.bboxes.append(item['box'])
            self.grasps.append(item['grasps'])
            self.objIDs.append(item['answer'])
            self.targets.append(item['target'])
            self.sentences.append(item['question'])
            self.semantics.append(item['program'])
            self.rgb_paths.append(os.path.join(seq_path, "rgb", im_name))
            self.depth_paths.append(os.path.join(seq_path, "depth", im_name))
            self.mask_paths.append(os.path.join(seq_path, "seg_mask_instances_combi", im_name))
            self.sent_indices.append(item['question_index'])
            self.sent_to_index[item['question_index']] = n
            n += 1
            
    def get_index_from_sent(self, sent_id):
        return self.sent_to_index[sent_id]

    def get_sent_from_index(self, n):
        return self.sent_indices[n]
    
    def _load_sent(self, sent_id):
        n = self.get_index_from_sent(sent_id)
        
        scene_id = self.scene_ids[n]
       
        img_path = os.path.join(self.root_dir, self.rgb_paths[n])
        img = self.get_image_from_path(img_path)
        
        x, y, w, h = self.bboxes[n]
        bbox = np.asarray([x, y, x+w, y+h])
        
        sent = self.sentences[n]
        
        target = self.targets[n]
        target_idx = self.class_instance_names[target]
        objID = self.objIDs[n]
        
        grasps = np.asarray(self.grasps[n])
        
        result = {'img': self.transform_img(img) if self.transform_img else img, 
                  'grasps':  self.transform_grasp(grasps, target_idx) if self.transform_grasp else None,
                  'grasp_rects': self.transform_grasp(grasps, target_idx) if self.transform_grasp else None,
                  'sentence': sent,
                  'target': target,
                  'objID': objID,
                  'bbox': bbox,
                  'target_idx': target_idx,
                  'sent_id': sent_id,
                  'scene_id': scene_id,
                  'img_path': img_path
                 }
        
        if self.with_depth:
            depth_path = os.path.join(self.root_dir, self.depth_paths[n])
            depth = self.get_depth_from_path(depth_path)
            result = {**result, 'depth': torch.from_numpy(depth) if self.transform_img else depth}

        if self.with_segm_mask:
            mask_path = os.path.join(self.root_dir, self.mask_paths[n])
            msk_full = self.get_mask_from_path(mask_path)
            msk = np.where(msk_full == objID, True, False)
            result = {**result, 'mask': torch.from_numpy(msk) if self.transform_img else msk}

        if self.with_grasp_masks:
            grasp_masks = self.transform_grasp.generate_masks(result['grasp_rects'])
            result = {**result, 'grasp_masks': grasp_masks}
        
        result = self.preprocess(result)
        
        return result
    
    def get_transform_mat(self, img_size, inverse=False):
        ori_h, ori_w = img_size
        inp_h, inp_w = self.input_size
        scale = min(inp_h / ori_h, inp_w / ori_w)
        new_h, new_w = ori_h * scale, ori_w * scale
        bias_x, bias_y = (inp_w - new_w) / 2., (inp_h - new_h) / 2.

        src = np.array([[0, 0], [ori_w, 0], [0, ori_h]], np.float32)
        dst = np.array([[bias_x, bias_y], [new_w + bias_x, bias_y],
                        [bias_x, new_h + bias_y]], np.float32)

        mat = cv2.getAffineTransform(src, dst)
        if inverse:
            mat_inv = cv2.getAffineTransform(dst, src)
            return mat, mat_inv
        return mat, None


    def preprocess(self, data):
        img = data["img"]
        sent = data["sentence"]
        if np.max(data["mask"]) <= 1.0:
            ins_mask = (data["mask"] * 255).astype(np.uint8)
        else:
            ins_mask = data["mask"]
        
        grasp_qua_mask = data["grasp_masks"]["qua"]
        grasp_ang_mask = data["grasp_masks"]["ang"]
        grasp_wid_mask = data["grasp_masks"]["wid"]

        img_size = img.shape[:2]
        mat, mat_inv = self.get_transform_mat(img_size, True)

        img = cv2.warpAffine(
            img, mat, self.input_size, flags=cv2.INTER_CUBIC,
            borderValue=[0.48145466 * 255, 0.4578275 * 255, 0.40821073 * 255]
        )

        img = torch.from_numpy(img.transpose((2, 0, 1)))
        if not isinstance(img, torch.FloatTensor):
            img = img.float()
        img.div_(255.).sub_(self.mean).div_(self.std)

        ins_mask = cv2.warpAffine(ins_mask,
                                    mat,
                                    self.input_size,
                                    flags=cv2.INTER_LINEAR,
                                    borderValue=0.)
        grasp_qua_mask = cv2.warpAffine(grasp_qua_mask,
                                mat,
                                self.input_size,
                                flags=cv2.INTER_LINEAR,
                                borderValue=0.)
        
        grasp_ang_mask = cv2.warpAffine(grasp_ang_mask,
                                mat,
                                self.input_size,
                                flags=cv2.INTER_LINEAR,
                                borderValue=0.)
        
        grasp_wid_mask = cv2.warpAffine(grasp_wid_mask,
                                mat,
                                self.input_size,
                                flags=cv2.INTER_LINEAR,
                                borderValue=0.)


        ins_mask = ins_mask / 255.
        grasp_qua_mask = grasp_qua_mask / 255.
        grasp_ang_mask = grasp_ang_mask * np.pi / 180.
        grasp_wid_mask = grasp_wid_mask / 255.
        grasp_sin_mask = np.sin(2 * grasp_ang_mask)
        grasp_cos_mask = np.cos(2 * grasp_ang_mask)

        word_vec = tokenize(sent, self.word_length, True).squeeze(0)

        data["img"] = img
        data["mask"] = ins_mask
        data["grasp_masks"]["qua"] = grasp_qua_mask
        data["grasp_masks"]["ang"] = grasp_ang_mask
        data["grasp_masks"]["wid"] = grasp_wid_mask
        data["grasp_masks"]["sin"] = grasp_sin_mask
        data["grasp_masks"]["cos"] = grasp_cos_mask
        data["word_vec"] = word_vec
        data["sentence"] = sent  # Keep the raw sentence
        data["inverse"] = mat_inv
        data["ori_size"] = np.array(img_size)
        
        # del data["sentence"]
        
        return data

    def __len__(self):
        return len(self.sent_indices)
    
    def __getitem__(self, n):
        sent_id = self.get_sent_from_index(n)
        data = self._load_sent(sent_id)
        
        return data
    
    @staticmethod
    def transform_grasp_inv(grasp_pt):
        pass
    
    # @functools.lru_cache(maxsize=None)
    def get_image_from_path(self, path):
        img_bgr = cv2.imread(path)
        if img_bgr is None:
            print(os.path.exists(path))
            print(path)
        img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        return img

    # @functools.lru_cache(maxsize=None)
    def get_mask_from_path(self, path):
        return cv2.imread(path, cv2.IMREAD_UNCHANGED)

    # @functools.lru_cache(maxsize=None)
    def get_depth_from_path(self, path):
        return cv2.imread(path, cv2.IMREAD_UNCHANGED).astype(np.float32) / 1000. # mm -> m
    
    def get_image(self, n):
        img_path = os.path.join(self.root_dir, self.imgs[n])
        return self.get_image_from_path(img_path)
    
    def get_annotated_image(self, n, text=True):
        sample = self.__getitem__(n)
        
        img, sent, grasps, bbox = sample['img'], sample['sentence'], sample['grasp_rects'], sample['bbox']
        if isinstance(img, torch.FloatTensor):
            img = img.permute(1,2,0)
            img = (img.cpu().numpy() * 255).astype(np.uint8)
        if self.transform_img:
            img = np.asarray(tfn.to_pil_image(img))
        if self.transform_grasp:
            #grasps = list(map(self.transform_grasp_inv, list(grasps)))
            grasps = self.transform_grasp.inverse(grasps)

        tmp = img.copy()
        for entry in grasps:
            ptA, ptB, ptC, ptD = [list(map(int, pt.tolist())) for pt in entry]
            tmp = cv2.line(tmp, ptA, ptB, (0,0,0xff), 2)
            tmp = cv2.line(tmp, ptD, ptC, (0,0,0xff), 2)
            tmp = cv2.line(tmp, ptB, ptC, (0xff,0,0), 2)
            tmp = cv2.line(tmp, ptA, ptD, (0xff,0,0), 2)
        
        tmp = cv2.rectangle(tmp, (bbox[0],bbox[1]), (bbox[2],bbox[3]), (0,255,0), 2)
        if text:
            tmp = cv2.putText(tmp, sent, (0,10), cv2.FONT_HERSHEY_PLAIN, 1, (0,0,0), 2, cv2.LINE_AA)
        return tmp

    def visualization(self, n, save_path):
        s = self.__getitem__(n)

        rgb = s['img']
        if isinstance(rgb, torch.FloatTensor):
            rgb = rgb.permute(1,2,0)
            rgb = (rgb.cpu().numpy() * 255).astype(np.uint8)
        depth = (0xff * s['depth'] / 3).astype(np.uint8)
        ii = self.get_annotated_image(n, text=False)
        sentence = s['sentence']
        msk = s['mask'].astype(np.uint8) / 255
        # msk_img = (rgb * 0.3).astype(np.uint8).copy()
        # msk_img[msk, 0] = 255

        fig = plt.figure(figsize=(25, 10))

        ax = fig.add_subplot(2, 4, 1)
        ax.imshow(rgb)
        ax.set_title('RGB')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 2)
        ax.imshow(depth, cmap='gray')
        ax.set_title('Depth')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 3)
        ax.imshow(msk)
        ax.set_title('Segm Mask')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 4)
        ax.imshow(ii)
        ax.set_title('Box & Grasp')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 5)
        plot = ax.imshow(s['grasp_masks']['qua'], cmap='jet', vmin=0, vmax=1)
        ax.set_title('Grasp quality')
        ax.axis('off')
        plt.colorbar(plot)

        ax = fig.add_subplot(2, 4, 6)
        plot = ax.imshow(s['grasp_masks']['sin'], cmap='rainbow', vmin=-1, vmax=1)
        ax.set_title('Angle-cosine')
        ax.axis('off')
        plt.colorbar(plot)

        ax = fig.add_subplot(2, 4, 7)
        plot = ax.imshow(s['grasp_masks']['cos'], cmap='rainbow', vmin=-1, vmax=1)
        ax.set_title('Angle-sine')
        ax.axis('off')
        plt.colorbar(plot)

        ax = fig.add_subplot(2, 4, 8)
        plot = ax.imshow(s['grasp_masks']['wid'], cmap='jet', vmin=0, vmax=1)
        ax.set_title('Width')
        ax.axis('off')
        plt.colorbar(plot)

        plt.suptitle(f"{sentence}", fontsize=20)
        plt.tight_layout()
        print("save")
        plt.savefig(os.path.join(save_path, f"sample_{n}.png"))
    
    @staticmethod
    def collate_fn(batch):
        return {
            "img": torch.stack([x["img"] for x in batch]),
            "depth": torch.stack([torch.from_numpy(x["depth"]).float() for x in batch]),
            "mask": torch.stack([torch.from_numpy(x["mask"]).float() for x in batch]),
            "grasp_masks" : {
                "qua": torch.stack([torch.from_numpy(x["grasp_masks"]["qua"]).float() for x in batch]),
                "sin": torch.stack([torch.from_numpy(x["grasp_masks"]["sin"]).float() for x in batch]),
                "cos": torch.stack([torch.from_numpy(x["grasp_masks"]["cos"]).float() for x in batch]),
                "wid": torch.stack([torch.from_numpy(x["grasp_masks"]["wid"]).float() for x in batch])
            },
            "word_vec": torch.stack([x["word_vec"].long() for x in batch]),
            "grasps": [x["grasps"] for x in batch],
            "target": [x["target"] for x in batch],
            "sentence": [x["sentence"] for x in batch],
            "bbox": [x["bbox"] for x in batch],
            "target_idx": [x["target_idx"] for x in batch],
            "sent_id": [x["sent_id"] for x in batch],
            "scene_id": [x["scene_id"] for x in batch],
            "inverse": [x["inverse"] for x in batch],
            "ori_size": [x["ori_size"] for x in batch],
            "img_path": [x["img_path"] for x in batch]
        }
            
        


class OCIDGraspDataset(Dataset):
    
    """ OCID-Grasp dataset """

    def __init__(self, 
                 cfg,
                 split):
        self.cfg = cfg
        self.split = split
        self.root_dir = cfg.root_dir
        self.img_size = cfg.img_size
        self.depth_factor = cfg.depth_factor
        self.with_grasp_masks = cfg.with_grasp_masks
        self.with_sem_masks = cfg.with_sem_masks
        self.with_ins_masks = cfg.with_ins_masks
        self.with_depth = cfg.with_depth
        self.grasp_transforms = GraspTransforms()

        aug_mode = "train" if self.split == "training_0" else "test"
        self.data_augmentor = DataAugmentor(cfg, mode=aug_mode)
        # self.data_augmentor = DataAugmentor(cfg)
        
        self._load_dicts()
        self.num_classes = len(cnames)

        with open(os.path.join(cfg.root_dir, "data_split", split + ".txt"), "r") as fid:
            self.meta = [x.strip().split(',') for x in fid.readlines()]


    def _load_dicts(self):
        cwd = os.getcwd()
        os.chdir(self.root_dir)
        from .OCID_sub_class_dict import cnames, colors, subnames, sub_to_class
        cnames_inv = {int(v):k for k,v in cnames.items()}
        subnames_inv = {v:k for k,v in subnames.items()}
        self.class_names = cnames 
        self.idx_to_class = cnames_inv
        self.class_instance_names = subnames
        self.idx_to_class_instance = subnames_inv
        self.instance_idx_to_class_idx = sub_to_class
        os.chdir(cwd)

    
    def _get_rgb_image(self, scene_id, img_f, data_dict):
        img_path = os.path.join(self.root_dir, scene_id, "rgb", img_f)
        img = cv2.imread(img_path, cv2.COLOR_BGR2RGB)
        data_dict["rgb"] = img

    
    def _get_depth_image(self, scene_id, img_f, data_dict):
        depth_path = os.path.join(self.root_dir, scene_id, "depth", img_f)
        depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED) / float(self.depth_factor)
        depth = 1 - (depth / np.max(depth))
        data_dict["depth"] = depth
        
    
    def _get_sem_mask(self, scene_id, img_f, data_dict):
        sem_mask = cv2.imread(os.path.join(self.root_dir, scene_id, "seg_mask_labeled_combi", img_f), cv2.IMREAD_UNCHANGED)
        data_dict["sem_mask"] = sem_mask
        return sem_mask

    
    def _get_ins_mask(self, scene_id, img_f, data_dict):
        # Load semantic mask first if the with_sem_masks set to False
        if "sem_mask" not in data_dict.keys():
            sem_mask = cv2.imread(os.path.join(self.root_dir, scene_id, "seg_mask_labeled_combi", img_f), cv2.IMREAD_UNCHANGED)
        else:
            sem_mask = data_dict["sem_mask"]
        ins_mask = cv2.imread(os.path.join(self.root_dir, scene_id, "seg_mask_instances_combi", img_f), cv2.IMREAD_UNCHANGED)

        labels     = []
        bboxes     = []
        ins_masks  = []

        props = regionprops(sem_mask)
        for prop in props:
            cls_id = prop.label
            
            # Get binary mask for each semantic class
            bin_mask = (sem_mask == cls_id).astype('int8')
            # Get corresponding semantic mask (may contains multiple instances)
            cls_ins_mask = (ins_mask * bin_mask)

            # Get regions for each instance
            ins_props = regionprops(cls_ins_mask)
            for ins in ins_props:
                labels.append(cls_id)
                bboxes.append([ins.bbox[1], ins.bbox[0], ins.bbox[3], ins.bbox[2], cls_id])
                mask = (cls_ins_mask == ins.label).astype('int8').astype('float32')
                ins_masks.append(mask)
        
        bboxes = np.array(bboxes).astype('float32')
        labels = np.array(labels)
        ins_masks  = np.array(ins_masks)

        data_dict["bboxes"] = bboxes
        data_dict["labels"] = labels
        data_dict["ins_masks"] = ins_masks


    
    def _get_per_cls_grasp_rects(self, scene_id, img_f, data_dict):
        anno_path = os.path.join(self.root_dir, scene_id, "Annotations_per_class", img_f[:-4])
        grasps_list = []
        for cls_id in os.listdir(anno_path):
            grasp_path = os.path.join(anno_path, cls_id, img_f[:-4]+".txt")
            with open(grasp_path, 'r') as f:
                points_list = []
                for count, line in enumerate(f):
                    line = line.rstrip()
                    [x, y] = line.split(' ')

                    x = float(x)
                    y = float(y)

                    pt = (x, y)
                    points_list.append(pt)

                    if len(points_list) == 4:
                        p1, p2, p3, p4 = points_list
                        center_x = (p1[0] + p3[0]) / 2
                        center_y = (p1[1] + p3[1]) / 2
                        width  = np.sqrt((p1[0] - p4[0]) * (p1[0] - p4[0]) + (p1[1] - p4[1]) * (p1[1] - p4[1]))
                        height = np.sqrt((p1[0] - p2[0]) * (p1[0] - p2[0]) + (p1[1] - p2[1]) * (p1[1] - p2[1]))
                        
                        # @NOTE
                        # Along x+ is 0 degree, increase by rotating anti-clockwise
                        # If you want to use opencv boxPoints & drawContours to visualize grasps
                        # Remember to take -theta as param :-)
                        theta = np.arctan2(p4[0] - p1[0], p4[1] - p1[1]) * 180 / np.pi
                        if theta > 0:
                            theta = theta-90
                        else:
                            theta = theta+90


                        grasps_list.append([center_x, center_y, width, height, theta, int(cls_id)])
                        points_list = []
        data_dict["raw_grasp_rects"] = grasps_list



    def _get_grasp_mask(self, scene_id, img_f, data_dict):
        grasp_rects = data_dict["raw_grasp_rects"]
        data_dict["grasp_masks"] = {}
        bboxes = data_dict["bboxes"]
        labels = data_dict["labels"]
        masks = data_dict["ins_masks"] # @NOTE there may exist instance with no grasp annotations, thus we need to filter them out
        assert bboxes.shape[0] == masks.shape[0], "inconsistent bounding boxes and instances, check the data"
        num_ins = bboxes.shape[0]

        ins_grasp_rects = []
        ins_grasp_masks = []
        ins_bboxes = []
        ins_masks = []
        ins_labels = []

        for i in range(num_ins):
            box = bboxes[i]
            mask = masks[i]
            label = labels[i]
            tmp = []
            for rect in grasp_rects:
                center_x, center_y, w, h = rect[:4]
                cls_id = rect[-1]
                # Grasp rect and bbox should have the same cls_id
                if int(cls_id) == int(box[4]):
                    # Center of grasp rect in bbox
                    if mask[int(center_y), int(center_x)]:
                        tmp.append(rect)
            if len(tmp) > 0:
                ins_grasp_masks.append(self.grasp_transforms.generate_masks(tmp))
                ins_grasp_rects.append(tmp)
                ins_bboxes.append(box)
                ins_masks.append(mask)
                ins_labels.append(label)

        data_dict["bboxes"] = np.asarray(ins_bboxes)
        data_dict["labels"] = np.asarray(ins_labels)
        data_dict["ins_masks"] = np.asarray(ins_masks)
        data_dict["ins_grasp_rects"] = ins_grasp_rects
        data_dict["grasp_masks"]["qua"] = np.asarray([gm["qua"] / 255 for gm in ins_grasp_masks])
        data_dict["grasp_masks"]["ang"] = np.asarray([gm["ang"] for gm in ins_grasp_masks])
        data_dict["grasp_masks"]["wid"] = np.asarray([gm["wid"] / 255 for gm in ins_grasp_masks])


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

    
    def __getitem__(self, index):
        data_dict = {}
        scene_id, img_f = self.meta[index]
        data_dict["scene_id"] = scene_id
        data_dict["img_f"] = img_f

        img_path = os.path.join(self.root_dir, scene_id, "rgb", img_f)
        img = cv2.imread(img_path)
        data_dict["rgb"] = img
        data_dict["ori_size"] = img.shape[:2]

        if self.with_depth:
            self._get_depth_image(scene_id, img_f, data_dict)
        if self.with_sem_masks:
            self._get_sem_mask(scene_id, img_f, data_dict)
        if self.with_ins_masks:
            self._get_ins_mask(scene_id, img_f, data_dict)
        if self.with_grasp_masks:
            self._get_per_cls_grasp_rects(scene_id, img_f, data_dict)
            self._get_grasp_mask(scene_id, img_f, data_dict)
        
        
        self.data_augmentor(data_dict)

        data_dict["grasp_masks"]["sin"] = np.sin(2 * data_dict["grasp_masks"]["ang"])
        data_dict["grasp_masks"]["cos"] = np.cos(2 * data_dict["grasp_masks"]["ang"])

        return data_dict


    def visualization(self, index, tgt_dir, with_preprocessing=False):
        data_dict = {}
        scene_id, img_f = self.meta[index]
        data_dict["scene_id"] = scene_id
        data_dict["img_f"] = img_f

        img_path = os.path.join(self.root_dir, scene_id, "rgb", img_f)
        img = cv2.imread(img_path, cv2.COLOR_BGR2RGB)
        data_dict["rgb"] = img

        if self.with_depth:
            self._get_depth_image(scene_id, img_f, data_dict)
        if self.with_sem_masks:
            self._get_sem_mask(scene_id, img_f, data_dict)
        if self.with_ins_masks:
            self._get_ins_mask(scene_id, img_f, data_dict)
        if self.with_grasp_masks:
            self._get_per_cls_grasp_rects(scene_id, img_f, data_dict)
            self._get_grasp_mask(scene_id, img_f, data_dict)
        if with_preprocessing:
            self.data_augmentor(data_dict)
            img = data_dict["rgb"].transpose((1,2,0))
        # img = img / 255.

        num_ins = data_dict["bboxes"].shape[0]
        
        fig = plt.figure(figsize=(25, 10))

        ax = fig.add_subplot(2, 4, 1)
        ax.imshow(np.clip(img[:, :, ::-1], 0.0, 1.0)*255)
        ax.set_title('RGB')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 2)
        ax.imshow(data_dict["depth"], cmap='gray')
        ax.set_title('Depth')
        ax.axis('off')

        ax = fig.add_subplot(2, 4, 3)
        ax.imshow(data_dict["sem_mask"])
        ax.set_title('Segm Mask')
        ax.axis('off')

        plt.tight_layout()
        plt.savefig(os.path.join(tgt_dir, f"raw-data.png"))
        plt.clf()
        plt.cla()
        plt.close()


        for i in range(num_ins):
            fig = plt.figure(figsize=(20, 2))

            ins_box = data_dict["bboxes"][i]
            ins_mask = data_dict["ins_masks"][i]
            ins_label = self.idx_to_class[data_dict["labels"][i]]
            ins_grasp_rects = data_dict["ins_grasp_rects"][i]
            ins_grasp_qua_masks = data_dict["grasp_masks"]["qua"][i]
            ins_grasp_ang_masks = data_dict["grasp_masks"]["ang"][i]
            ins_grasp_wid_masks = data_dict["grasp_masks"]["wid"][i]

            grasps = self.grasp_transforms.inverse(ins_grasp_rects)
            tmp = img.copy()
            h, w, c = tmp.shape
            for entry in grasps:
                ptA, ptB, ptC, ptD = [list(map(int, pt.tolist())) for pt in entry]
                tmp = cv2.line(tmp, ptA, ptB, (0,0,0xff), 2)
                tmp = cv2.line(tmp, ptD, ptC, (0,0,0xff), 2)
                tmp = cv2.line(tmp, ptB, ptC, (0xff,0,0), 2)
                tmp = cv2.line(tmp, ptA, ptD, (0xff,0,0), 2)
            if ins_box[0] <= 1:
                tmp = cv2.rectangle(tmp, (int(ins_box[0]*w),int(ins_box[1]*h)), (int(ins_box[2]*w),int(ins_box[3]*h)), (0,0,0xff), 2)
            else:
                tmp = cv2.rectangle(tmp, (int(ins_box[0]),int(ins_box[1])), (int(ins_box[2]),int(ins_box[3])), (0,0,0xff), 2)

            ax = fig.add_subplot(1, 5, 1)
            ax.imshow(tmp[:, :, ::-1])
            ax.set_title('Bboxes & Grasps')
            ax.axis('off')

            ax = fig.add_subplot(1, 5, 2)
            tmp_mask = np.expand_dims(ins_mask, axis=-1).repeat(3, axis=-1)
            ax.imshow(tmp_mask*0.6 + tmp[:, :, ::-1] *0.4)
            ax.set_title('ins mask')
            ax.axis('off')

            ax = fig.add_subplot(1, 5, 3)
            plot = ax.imshow(ins_grasp_qua_masks, cmap='jet', vmin=0, vmax=1)
            ax.set_title('Grasp quality')
            ax.axis('off')
            plt.colorbar(plot)

            ax = fig.add_subplot(1, 5, 4)
            plot = ax.imshow(ins_grasp_ang_masks, cmap='rainbow', vmin=-1, vmax=1)
            ax.set_title('Grasp angle')
            ax.axis('off')
            plt.colorbar(plot)

            ax = fig.add_subplot(1, 5, 5)
            plot = ax.imshow(ins_grasp_wid_masks, cmap='rainbow', vmin=0, vmax=1)
            ax.set_title('Grasp width')
            ax.axis('off')
            plt.colorbar(plot)

            plt.tight_layout()
            plt.savefig(os.path.join(tgt_dir, f"ins-{i}-{ins_label}.png"))


    @staticmethod
    def collate_fn(batch):
        return {
            "scene_id": [x["scene_id"] for x in batch],
            "img_f": [x["img_f"] for x in batch],
            "ori_size": batch[0]["ori_size"],
            "rgb": torch.stack([torch.from_numpy(x["rgb"]) for x in batch]),
            "depth": torch.stack([torch.from_numpy(x["depth"]) for x in batch]).unsqueeze(1),
            "labels": [torch.from_numpy(x["labels"]).long() for x in batch],
            "bboxes": [torch.from_numpy(x["bboxes"]) for x in batch],
            "ins_masks": [torch.from_numpy(x["ins_masks"]).float() for x in batch],
            "sem_mask": torch.stack([torch.from_numpy(x["sem_mask"]).float() for x in batch]),
            "grasp_rects": [x["ins_grasp_rects"] for x in batch],
            "grasp_masks" : {
                "qua": [torch.from_numpy(x["grasp_masks"]["qua"]).float() for x in batch],
                "sin": [torch.from_numpy(x["grasp_masks"]["sin"]).float() for x in batch],
                "cos": [torch.from_numpy(x["grasp_masks"]["cos"]).float() for x in batch],
                "wid": [torch.from_numpy(x["grasp_masks"]["wid"]).float() for x in batch]
            },
        }