import os import io import re import json import pickle from typing import Optional import cv2 import numpy as np import torch import webdataset as wds from torch.utils.data import DataLoader class GraspAnythingPPWebDataset: """ WebDataset version for GraspAnything++ (no split filtering). Expected each sample in tar contains: { "__key__": sample_id, "jpg": , # RGB image bytes "pt": , # grasp tensor .pt bytes "pkl": , # instruction pickle bytes "json": , # metadata json bytes } Output format is aligned with the original GraspAnythingPPDatasetwomask: { "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, tar_pattern, transform_img=None, transform_grasp=None, input_size=416, word_length=20, with_depth=False, with_grasp_masks=True, width_factor=100, shardshuffle=True, sampleshuffle=1000, repeat=False, nodesplitter=None, ): self.tar_pattern = tar_pattern self.with_depth = bool(with_depth) 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) # build dataset pipeline dataset = wds.WebDataset( self.tar_pattern, shardshuffle=shardshuffle, nodesplitter=nodesplitter, ) if repeat: dataset = dataset.repeat() if sampleshuffle and sampleshuffle > 0: dataset = dataset.shuffle(sampleshuffle) dataset = dataset.map(self._decode_and_process_sample) self.dataset = dataset # --------------------------------------------------------- # Public # --------------------------------------------------------- def get_dataset(self): return self.dataset def get_dataloader( self, batch_size, num_workers=4, pin_memory=True, persistent_workers=False, drop_last=False, ): return DataLoader( self.dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, persistent_workers=persistent_workers if num_workers > 0 else False, drop_last=drop_last, collate_fn=self.collate_fn, ) # --------------------------------------------------------- # 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 not isinstance(instruction, str): raise TypeError(f"Expected instruction to be a string, got {type(instruction)}") return instruction def _load_ga_grasps_from_bytes(self, pt_bytes): 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 (same logic as your original dataset) # --------------------------------------------------------- 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_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 # IMPORTANT: # You must already have tokenize(...) in your project, # same as your original dataset. 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): """ Convert one WebDataset sample to the same dict format as the original dataset. """ if "jpg" not in sample: raise KeyError(f"Sample missing 'jpg'. Keys: {list(sample.keys())}") if "pt" not in sample: raise KeyError(f"Sample missing 'pt'. Keys: {list(sample.keys())}") if "pkl" not in sample: raise KeyError(f"Sample missing 'pkl'. 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"]) grasps_ga = self._load_ga_grasps_from_bytes(sample["pt"]) sentence = self._load_instruction_from_bytes(sample["pkl"]) 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"), "sample_id": sample_id, } if self.with_depth: # only valid if you packed depth into shards too 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: result["grasp_masks"] = self.transform_grasp.generate_masks(result["grasp_rects"]) return self.preprocess(result) # --------------------------------------------------------- # Collate (kept compatible with your original dataset) # --------------------------------------------------------- @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], "inverse": [x["inverse"] for x in batch], "ori_size": [x["ori_size"] 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