def __getitem__()

in experiments/grasp_stability/train.py [0:0]

• 35 lines of code
• 11 McCabe index (conditional complexity)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        fileID = idx // self.numGroup
        if fileID != self.dataFileID:
            self.dataList = self.load_data(fileID)
            self.dataFileID = fileID

        sample = {}

        data = self.dataList

        for k in self.fields:
            d = data[k][idx % self.numGroup]

            if k in ["tactileColorL", "tactileColorR", "visionColor"]:
                d = d[:, :, :3]
                # print(k, d.min(), d.max())

            if k in ["tactileDepthL", "tactileDepthR", "visionDepth"]:
                d = np.dstack([d, d, d])

            if k in ["tactileDepthL", "tactileDepthR"]:
                d = d / 0.002 * 255
                d = np.clip(d, 0, 255).astype(np.uint8)
                # print("depth min", d.min(), "max", d.max())

            if k in ["visionDepth"]:
                d = (d * 255).astype(np.uint8)

            if k in [
                "tactileColorL",
                "tactileColorR",
                "visionColor",
                "visionDepth",
            ]:
                if self.transform:
                    d = self.transform(d)

            if k in [
                "tactileDepthL",
                "tactileDepthR",
            ]:
                # print("before", d.min(), d.max(), d.mean(), d.std())
                d = self.transformDepth(d)
                # d = (d + 2) / 0.05
                # print("after", d.min(), d.max(), d.mean(), d.std())

            sample[k] = d

        return sample