import os from typing import List, Union import json import itertools import cv2 import lmdb import numpy as np import pyarrow as pa import torch import re import glob import pickle 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 import webdataset as wds import io import zlib from torch.utils.data import get_worker_info 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 DDPNodeSplitter: def __init__(self, rank=None, world_size=None): self.rank = rank self.world_size = world_size def __call__(self, src): local_rank = self.rank local_world_size = self.world_size if local_rank is None or local_world_size is None: if torch.distributed.is_available() and torch.distributed.is_initialized(): local_rank = torch.distributed.get_rank() local_world_size = torch.distributed.get_world_size() else: local_rank = 0 local_world_size = 1 if local_world_size <= 1: yield from src return yield from itertools.islice(src, local_rank, None, local_world_size) class DDPNodeSplitterWithSharedTail(DDPNodeSplitter): def __init__(self, shared_tail_url, rank=None, world_size=None): super().__init__(rank=rank, world_size=world_size) self.shared_tail_url = os.path.normpath(shared_tail_url) def __call__(self, src): local_rank = self.rank local_world_size = self.world_size if local_rank is None or local_world_size is None: if torch.distributed.is_available() and torch.distributed.is_initialized(): local_rank = torch.distributed.get_rank() local_world_size = torch.distributed.get_world_size() else: local_rank = 0 local_world_size = 1 shared_tail = None regular_index = 0 for shard in src: shard_url = os.path.normpath(shard["url"]) if shard_url == self.shared_tail_url: shared_tail = shard continue if regular_index % local_world_size == local_rank: yield shard regular_index += 1 if shared_tail is not None: yield shared_tail class SharedTailSampleSelector: def __init__(self, shared_tail_url, rank, samples_per_rank): self.shared_tail_name = os.path.basename(shared_tail_url) self.start = int(rank) * int(samples_per_rank) self.stop = self.start + int(samples_per_rank) self.index = 0 def __call__(self, sample): sample_url = os.path.basename(str(sample.get("__url__", ""))) if sample_url != self.shared_tail_name: return True index = self.index self.index += 1 return self.start <= index < self.stop def wds_worker_splitter(src): worker_info = get_worker_info() if worker_info is None: yield from src return yield from itertools.islice(src, worker_info.id, None, worker_info.num_workers) 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"]) 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] }, } class GraspTransformsGAPlusPlus: """ For GA++ grasp format: (N, 6) = [score, cx, cy, angle_deg, width, length] Output (N, 6) for OCID-style pipeline: [center_x, center_y, w_rect, h_rect, theta, target] where theta is in (-90, 90] degrees (OCID internal convention). generate_masks() outputs the same dict keys as OCID: {'pos': uint8, 'qua': uint8, 'ang': uint8, 'wid': uint8} """ def __init__(self, width_factor=100, width=416, height=416, mask_mode="legacy", gaussian_sigma=3): self.width_factor = float(width_factor) self.width = int(width) self.height = int(height) self.mask_mode = mask_mode self.gaussian_sigma = gaussian_sigma @staticmethod def _angle180_to_theta(angle_deg): """ Convert angle in [0,180) to theta in (-90, 90] (OCID internal). """ a = float(angle_deg) return a - 180.0 if a > 90.0 else a def __call__(self, grasps_ga, target=0): """ :param grasps_ga: torch.Tensor or np.ndarray with shape (N,6) [score, cx, cy, angle_deg, width, length] :return: np.ndarray (N,6) float32 [cx, cy, w_rect, h_rect, theta, target] """ if grasps_ga is None: return np.zeros((0, 6), dtype=np.float32) if hasattr(grasps_ga, "detach"): # torch tensor g = grasps_ga.detach().cpu().numpy() else: g = np.asarray(grasps_ga) if g.ndim != 2 or g.shape[1] < 6: raise ValueError(f"Expected GA++ grasps shape (N,6), got {g.shape}") cx = g[:, 1].astype(np.float32) cy = g[:, 2].astype(np.float32) angle_deg = g[:, 3] width = g[:, 4].astype(np.float32) length = g[:, 5].astype(np.float32) theta = np.array([self._angle180_to_theta(a) for a in angle_deg], dtype=np.float32) height = length tgt = np.full_like(cx, float(target), dtype=np.float32) out = np.stack([cx, cy, width, height, theta, tgt], axis=1).astype(np.float32) return out def inverse(self, grasp_rectangles): """ Convert [cx,cy,w,h,theta] back to 4-corner polygon (x,y) for visualization. """ boxes = [] for rect in grasp_rectangles: center_x, center_y, width, height, theta = rect[:5] box = ((center_x, center_y), (width, height), -(theta + 180.0)) pts = cv2.boxPoints(box) boxes.append(np.intp(pts)) return boxes def _rect_pixels(self, center_x, center_y, w_rect, h_rect, theta, width=None, height=None): canvas_width = self.width if width is None else int(width) canvas_height = self.height if height is None else int(height) r_rect = ((float(center_x), float(center_y)), (float(w_rect) / 2, float(h_rect)), -(float(theta) + 180.0)) box = cv2.boxPoints(r_rect) box = np.intp(box) rr, cc = polygon(box[:, 0], box[:, 1]) valid = (rr >= 0) & (rr < canvas_width) & (cc >= 0) & (cc < canvas_height) return rr[valid], cc[valid] def _generate_masks_quality_max(self, grasp_rectangles, width=None, height=None): canvas_width = self.width if width is None else int(width) canvas_height = self.height if height is None else int(height) sigma = getattr(self, "gaussian_sigma", 3) pos_out = np.zeros((canvas_height, canvas_width), dtype=np.float32) qua_out = np.zeros((canvas_height, canvas_width), dtype=np.float32) owner_qua = np.zeros((canvas_height, canvas_width), dtype=np.float32) ang_out = np.zeros((canvas_height, canvas_width), dtype=np.float32) wid_out = np.zeros((canvas_height, canvas_width), dtype=np.float32) for rect in grasp_rectangles: center_x, center_y, w_rect, h_rect, theta = rect[:5] rr, cc = self._rect_pixels( center_x, center_y, w_rect, h_rect, theta, width=canvas_width, height=canvas_height, ) if rr.size == 0: continue rect_pos = np.zeros((canvas_height, canvas_width), dtype=np.float32) rect_pos[cc, rr] = 1.0 rect_qua = gaussian(rect_pos, sigma, preserve_range=True) update = (rect_pos > 0) & (rect_qua > owner_qua) pos_out[cc, rr] = 1.0 qua_out = np.maximum(qua_out, rect_qua) owner_qua[update] = rect_qua[update] if theta < 0: angle = int(theta + 180) else: angle = int(theta) width = np.clip(w_rect, 0.0, self.width_factor) / self.width_factor ang_out[update] = angle wid_out[update] = width pos_u8 = (pos_out * 255).astype(np.uint8) qua_u8 = (qua_out * 255).astype(np.uint8) ang_u8 = ang_out.astype(np.uint8) wid_u8 = (gaussian(wid_out, sigma, preserve_range=True) * 255).astype(np.uint8) return {"pos": pos_u8, "qua": qua_u8, "ang": ang_u8, "wid": wid_u8} def generate_masks(self, grasp_rectangles, width=None, height=None): """ Same outputs as OCID: {'pos': uint8, 'qua': uint8, 'ang': uint8, 'wid': uint8} """ canvas_width = self.width if width is None else int(width) canvas_height = self.height if height is None else int(height) mask_mode = getattr(self, "mask_mode", "legacy") sigma = getattr(self, "gaussian_sigma", 3) if mask_mode in ("quality_max", "max_quality", "per_grasp_max"): return self._generate_masks_quality_max( grasp_rectangles, width=canvas_width, height=canvas_height ) pos_out = np.zeros((canvas_height, canvas_width), dtype=np.float32) ang_out = np.zeros((canvas_height, canvas_width), dtype=np.float32) wid_out = np.zeros((canvas_height, canvas_width), dtype=np.float32) for rect in grasp_rectangles: center_x, center_y, w_rect, h_rect, theta = rect[:5] rr, cc = self._rect_pixels( center_x, center_y, w_rect, h_rect, theta, width=canvas_width, height=canvas_height, ) if rr.size == 0: continue pos_out[cc, rr] = 1.0 # Store angle map in [0,180) if theta < 0: ang_out[cc, rr] = int(theta + 180) else: ang_out[cc, rr] = int(theta) # Width map normalized by width_factor (same as OCID) wid_out[cc, rr] = np.clip(w_rect, 0.0, self.width_factor) / self.width_factor qua_out = (gaussian(pos_out, sigma, preserve_range=True) * 255).astype(np.uint8) pos_u8 = (pos_out * 255).astype(np.uint8) ang_u8 = ang_out.astype(np.uint8) wid_u8 = (gaussian(wid_out, sigma, preserve_range=True) * 255).astype(np.uint8) return {"pos": pos_u8, "qua": qua_out, "ang": ang_u8, "wid": wid_u8} import io import json import pickle import cv2 import numpy as np import torch import webdataset as wds class GraspAnythingPPWebDataset: """ WebDataset version for GraspAnything++. 每个样本期望包含： { "__key__": sample_id, "jpg": , "pt"/"ptz": , "pkl"/"pklz": , optional "npy"/"npyz": , "json": , } 输出与原始 dataset 保持兼容： { "img": Tensor[C,H,W], "word_vec": Tensor[L], "grasps": ..., "grasp_rects": ..., "sentence": str, "obj_id": int, "part_id": int, "grasp_path": str, "img_path": str, "sample_id": str, "inverse": np.ndarray, "ori_size": np.ndarray, optional: "depth" "grasp_masks" } """ def __init__( self, transform_img=None, transform_grasp=None, input_size=416, word_length=20, with_depth=False, with_segm_mask=False, with_grasp_masks=True, ): self.with_depth = bool(with_depth) self.with_segm_mask = bool(with_segm_mask) self.with_grasp_masks = bool(with_grasp_masks) self.input_size = (int(input_size), int(input_size)) self.word_length = int(word_length) self.transform_img = transform_img self.transform_grasp = transform_grasp if self.transform_grasp is None: raise ValueError("transform_grasp must be provided.") 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) def make_train_dataset( self, tar_files, batch_size, steps_per_epoch, sampleshuffle=1000, rank=None, world_size=None, shared_tail_tar_file=None, shared_tail_samples_per_rank=0, handler=None, ): """ 构建训练用 WebDataset 数据管线。 每个逻辑 epoch 读取一个不重叠的 shard 分块。最后一个不完整 shard 由所有 DDP rank 共同读取，再按样本范围无重叠切分。 """ if handler is None: handler = wds.handlers.reraise_exception if shared_tail_tar_file is None: node_splitter = DDPNodeSplitter(rank=rank, world_size=world_size) else: node_splitter = DDPNodeSplitterWithSharedTail( shared_tail_tar_file, rank=rank, world_size=world_size ) dataset = wds.WebDataset( tar_files, resampled=False, shardshuffle=False, nodesplitter=node_splitter, workersplitter=wds_worker_splitter, handler=handler, ) if shared_tail_tar_file is not None: dataset = dataset.select( SharedTailSampleSelector( shared_tail_tar_file, rank, shared_tail_samples_per_rank ) ) if sampleshuffle and sampleshuffle > 0: dataset = dataset.shuffle(sampleshuffle) dataset = dataset.map(self._decode_and_process_sample, handler=handler) # 在 WebDataset 管线内部完成 batch，而不是交给 DataLoader 自动 batch dataset = dataset.batched( batch_size, partial=False, collation_fn=self.collate_fn, ) # 明确规定每个 epoch 的 step 数 dataset = dataset.with_epoch(steps_per_epoch) return dataset def make_eval_dataset( self, tar_files, batch_size, steps_per_epoch=None, rank=None, world_size=None, handler=None, ): if handler is None: handler = wds.handlers.reraise_exception dataset = wds.WebDataset( tar_files, resampled=False, shardshuffle=False, nodesplitter=DDPNodeSplitter(rank=rank, world_size=world_size), workersplitter=wds_worker_splitter, handler=handler, ) dataset = dataset.map(self._decode_and_process_sample, handler=handler) dataset = dataset.batched( batch_size, partial=True, collation_fn=self.collate_fn, ) if steps_per_epoch is not None: dataset = dataset.with_epoch(steps_per_epoch) return dataset def make_val_dataset(self, *args, **kwargs): return self.make_eval_dataset(*args, **kwargs) def make_test_dataset(self, *args, **kwargs): return self.make_eval_dataset(*args, **kwargs) # --------------------------------------------------------- # Byte decoders # --------------------------------------------------------- def _load_image_from_bytes(self, image_bytes): img_arr = np.frombuffer(image_bytes, dtype=np.uint8) img_bgr = cv2.imdecode(img_arr, cv2.IMREAD_COLOR) if img_bgr is None: raise ValueError("Failed to decode image bytes.") img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) return img def _load_depth_from_bytes(self, depth_bytes): depth_arr = np.frombuffer(depth_bytes, dtype=np.uint8) depth = cv2.imdecode(depth_arr, cv2.IMREAD_UNCHANGED) if depth is None: raise ValueError("Failed to decode depth bytes.") return depth.astype(np.float32) / 1000.0 def _load_instruction_from_bytes(self, pkl_bytes): instruction = pickle.loads(pkl_bytes) if isinstance(instruction, str): return instruction if isinstance(instruction, dict): for key in ("instruction", "sentence", "text", "prompt", "question", "description"): value = instruction.get(key) if isinstance(value, str): return value if isinstance(instruction, (list, tuple)): for value in instruction: if isinstance(value, str): return value raise TypeError(f"Expected instruction text in pickle, got {type(instruction)}") def _load_part_mask_from_bytes(self, npy_bytes): mask = np.load(io.BytesIO(npy_bytes), allow_pickle=False) mask = np.asarray(mask) mask = np.squeeze(mask) if mask.ndim != 2: raise ValueError(f"Expected a 2D GA++ part mask, got shape {mask.shape}") return (mask > 0).astype(np.uint8) * 255 def _load_ga_grasps_from_bytes(self, pt_bytes): try: x = torch.load(io.BytesIO(pt_bytes), map_location="cpu", weights_only=True) except TypeError: # PyTorch < 2.0 has no weights_only argument. x = torch.load(io.BytesIO(pt_bytes), map_location="cpu") if torch.is_tensor(x): return x if isinstance(x, dict): for k in ["grasps", "grasp", "positive_grasps", "data"]: if k in x and torch.is_tensor(x[k]): return x[k] raise ValueError(f"Unsupported .pt content type for grasps: {type(x)}") def _load_meta_from_bytes(self, json_bytes): if isinstance(json_bytes, bytes): return json.loads(json_bytes.decode("utf-8")) if isinstance(json_bytes, str): return json.loads(json_bytes) if isinstance(json_bytes, dict): return json_bytes raise TypeError(f"Unsupported metadata type: {type(json_bytes)}") # --------------------------------------------------------- # Preprocess # --------------------------------------------------------- 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.0, (inp_h - new_h) / 2.0 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"] 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.0).sub_(self.mean).div_(self.std) data["img"] = img data["inverse"] = mat_inv data["ori_size"] = np.array(img_size) if self.with_segm_mask: part_mask = cv2.warpAffine( data["mask"], mat, self.input_size, flags=cv2.INTER_NEAREST, borderValue=0.0, ) data["mask"] = (part_mask > 0).astype(np.float32) if self.with_depth and "depth" in data: depth = data["depth"] depth = cv2.warpAffine( depth, mat, self.input_size, flags=cv2.INTER_NEAREST, borderValue=0.0, ) data["depth"] = depth if self.with_grasp_masks: qua = data["grasp_masks"]["qua"] ang = data["grasp_masks"]["ang"] wid = data["grasp_masks"]["wid"] qua = cv2.warpAffine(qua, mat, self.input_size, flags=cv2.INTER_LINEAR, borderValue=0.0) ang = cv2.warpAffine(ang, mat, self.input_size, flags=cv2.INTER_LINEAR, borderValue=0.0) wid = cv2.warpAffine(wid, mat, self.input_size, flags=cv2.INTER_LINEAR, borderValue=0.0) qua = qua / 255.0 wid = wid / 255.0 ang = ang * np.pi / 180.0 sin = np.sin(2 * ang) cos = np.cos(2 * ang) data["grasp_masks"]["qua"] = qua data["grasp_masks"]["ang"] = ang data["grasp_masks"]["wid"] = wid data["grasp_masks"]["sin"] = sin data["grasp_masks"]["cos"] = cos # 这里沿用你项目里的 tokenize word_vec = tokenize(sent, self.word_length, True).squeeze(0) data["word_vec"] = word_vec data["sentence"] = sent return data # --------------------------------------------------------- # Sample processing # --------------------------------------------------------- def _decode_and_process_sample(self, sample): if "jpg" not in sample: raise KeyError(f"Sample missing 'jpg'. Keys: {list(sample.keys())}") if "pt" not in sample and "ptz" not in sample: raise KeyError(f"Sample missing 'pt'/'ptz'. Keys: {list(sample.keys())}") if "pkl" not in sample and "pklz" not in sample: raise KeyError(f"Sample missing 'pkl'/'pklz'. Keys: {list(sample.keys())}") if self.with_segm_mask and "npy" not in sample and "npyz" not in sample: raise KeyError(f"Sample missing 'npy'/'npyz' part mask. Keys: {list(sample.keys())}") if "json" not in sample: raise KeyError(f"Sample missing 'json'. Keys: {list(sample.keys())}") meta = self._load_meta_from_bytes(sample["json"]) img = self._load_image_from_bytes(sample["jpg"]) pt_bytes = sample["pt"] if "pt" in sample else zlib.decompress(sample["ptz"]) pkl_bytes = sample["pkl"] if "pkl" in sample else zlib.decompress(sample["pklz"]) grasps_ga = self._load_ga_grasps_from_bytes(pt_bytes) sentence = self._load_instruction_from_bytes(pkl_bytes) grasp_rects = self.transform_grasp(grasps_ga, target=0) if self.transform_grasp else None sample_id = meta.get("sample_id", sample.get("__key__", "")) prompt_id = meta.get("prompt_id", "") obj_id = int(meta.get("obj_id", -1)) part_id = int(meta.get("part_id", -1)) result = { "img": self.transform_img(img) if self.transform_img else img, "grasps": grasp_rects, "grasp_rects": grasp_rects, "sentence": sentence, "obj_id": obj_id, "part_id": part_id, "grasp_path": meta.get("grasp_file", f"{sample_id}.pt"), "img_path": meta.get("image_file", f"{prompt_id}.jpg"), "mask_path": meta.get("mask_file", ""), "sample_id": sample_id, } if self.with_segm_mask: npy_bytes = sample["npy"] if "npy" in sample else zlib.decompress(sample["npyz"]) result["mask"] = self._load_part_mask_from_bytes(npy_bytes) if self.with_depth: if "depth" in sample: result["depth"] = self._load_depth_from_bytes(sample["depth"]) else: raise KeyError("with_depth=True but sample does not contain 'depth' field.") if self.with_grasp_masks: orig_h, orig_w = img.shape[:2] result["grasp_masks"] = self.transform_grasp.generate_masks( result["grasp_rects"], width=orig_w, height=orig_h, ) return self.preprocess(result) # --------------------------------------------------------- # Collate # --------------------------------------------------------- @staticmethod def collate_fn(batch): out = { "img": torch.stack([x["img"] for x in batch]), "word_vec": torch.stack([x["word_vec"].long() for x in batch]), "grasps": [x["grasps"] for x in batch], "sentence": [x["sentence"] for x in batch], "obj_id": [x.get("obj_id", -1) for x in batch], "part_id": [x.get("part_id", -1) for x in batch], "sample_id": [x.get("sample_id", "") for x in batch], "grasp_path": [x.get("grasp_path", "") for x in batch], "img_path": [x.get("img_path", "") for x in batch], "mask_path": [x.get("mask_path", "") for x in batch], "inverse": [x["inverse"] for x in batch], "ori_size": [x["ori_size"] for x in batch], } if "mask" in batch[0]: out["mask"] = torch.stack( [torch.from_numpy(x["mask"]).float() for x in batch] ) if "depth" in batch[0]: out["depth"] = torch.stack([torch.from_numpy(x["depth"]).float() for x in batch]) if "grasp_masks" in batch[0]: out["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]), } return out