| 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": <bytes>, # RGB image bytes |
| "pt": <bytes>, # grasp tensor .pt bytes |
| "pkl": <bytes>, # instruction pickle bytes |
| "json": <bytes>, # 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 |
