pointcept/datasets/scanobjectnn.py

"""
ScanObjectNN Dataset

Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import os
import glob
import h5py
import numpy as np
import torch
from collections.abc import Sequence
import open3d as o3d
import copy

from pointcept.utils.logger import get_root_logger

from .builder import DATASETS
from .defaults import DefaultDataset
from .transform import Compose


@DATASETS.register_module()
class ScanObjectNNRawDataset(DefaultDataset):
    """
    ScanObjectNN Dataset loaded from a standard folder structure.
    """

    def __init__(self, if_img=True, crop_h=518, crop_w=518, patch_size=14, **kwargs):
        self.crop_h = crop_h
        self.crop_w = crop_w
        self.patch_size = patch_size
        self.patch_h = crop_h // patch_size
        self.patch_w = crop_w // patch_size
        self.if_img = if_img
        super().__init__(**kwargs)

    def get_data_list(self):
        """
        Overrides the parent method.
        Scans the data directory (or directories) to find all relevant .bin files
        for the specified split(s).
        """
        if isinstance(self.split, str):
            split_list = [self.split]
        elif isinstance(self.split, Sequence):
            split_list = self.split
        else:
            raise NotImplementedError

        data_list = []
        logger = get_root_logger()

        for split in split_list:
            split_path = os.path.join(self.data_root, split)
            if not os.path.isdir(split_path):
                raise FileNotFoundError(f"Split directory not found at: {split_path}")

            logger.info(f"Scanning for .bin files in {split_path}")
            # Use glob to find all .bin files within the class subdirectories
            all_bin_files = glob.glob(os.path.join(split_path, "*", "*.bin"))

            # Filter out auxiliary files, keeping only the main object files
            filtered_files = [
                path
                for path in all_bin_files
                if not path.endswith("_part.bin") and not path.endswith("_indices.bin")
            ]

            if not filtered_files:
                logger.warning(
                    f"No valid .bin files found in {split_path}. Please check the file structure."
                )

            data_list.extend(filtered_files)

        return sorted(data_list)  # Sort for reproducibility

    def get_data(self, idx):
        """
        Overrides the parent method.
        Loads and parses a single .bin file based on its path.
        The logic for parsing the binary file format remains the same.
        """
        data_path = self.data_list[idx % len(self.data_list)]

        try:
            raw_data = np.fromfile(data_path, dtype=np.float32)
        except IOError as e:
            logger = get_root_logger()
            logger.error(f"Could not read file {data_path}: {e}")
            # Skip to the next sample if the current one is unreadable
            return self.get_data((idx + 1) % len(self.data_list))

        if raw_data.size == 0:
            logger = get_root_logger()
            logger.warning(f"File is empty, skipping: {data_path}")
            return self.get_data((idx + 1) % len(self.data_list))

        # First float is the number of points
        num_points = int(raw_data[0])
        point_data_flat = raw_data[1:]

        # Validate that the file is not corrupt
        expected_size = num_points * 11
        if point_data_flat.size != expected_size:
            logger = get_root_logger()
            logger.warning(
                f"Data corruption detected in {data_path}: "
                f"Expected {expected_size} floats for {num_points} points, but found {point_data_flat.size}. Skipping."
            )
            return self.get_data((idx + 1) % len(self.data_list))

        # Reshape into a (num_points, 11) array
        point_cloud = point_data_flat.reshape(num_points, 11)
        rot_mat = np.array(
            [
                [1.0, 0.0, 0.0],
                [0.0, 0.0, -1.0],
                [0.0, 1.0, 0.0],
            ],
            dtype=np.float32,
        )
        rotated_coord = point_cloud[:, :3].astype(np.float32) @ rot_mat.T
        rotated_normal = point_cloud[:, 3:6].astype(np.float32) @ rot_mat.T
        if self.if_img:
            imgs_tensor = torch.empty(
                (0, 3, self.patch_h * self.patch_size, self.patch_w * self.patch_size)
            )
            img_num = np.array([imgs_tensor.shape[0]], dtype=np.int32)
            correspondence_infos = np.ones(
                (point_cloud.shape[0], imgs_tensor.shape[0], 2),
                dtype=np.float32,
                # (data_dict["coord"].shape[0], len(correspondences_path), 2), dtype=np.int32
            ) * (-1)
            # Pack the data into the dictionary format expected by the framework
            # Attributes order: x, y, z, nx, ny, nz, r, g, b, instance_label, semantic_label
            data_dict = {
                "coord": rotated_coord.astype(np.float32),
                "normal": rotated_normal.astype(np.float32),
                "color": point_cloud[:, 6:9].astype(np.float32),
                "instance": point_cloud[:, 9].astype(np.int64),
                "segment": point_cloud[:, 10].astype(np.int64),
                "name": self.get_data_name(idx),
                "split": self.get_split_name(idx),
                "images": imgs_tensor,
                "img_num": img_num,
                "correspondence": correspondence_infos,
            }
        else:
            data_dict = {
                "coord": rotated_coord.astype(np.float32),
                "normal": rotated_normal.astype(np.float32),
                "color": point_cloud[:, 6:9].astype(np.float32),
                "instance": point_cloud[:, 9].astype(np.int64),
                "segment": point_cloud[:, 10].astype(np.int64),
                "name": self.get_data_name(idx),
                "split": self.get_split_name(idx),
            }
        return data_dict

    def get_data_name(self, idx):
        """
        Overrides the parent method.
        Returns the filename without extension as the data name.
        Example: '005_00001'
        """
        path = self.data_list[idx % len(self.data_list)]
        return os.path.splitext(os.path.basename(path))[0]

    def get_split_name(self, idx):
        """
        Overrides the parent method.
        Correctly extracts the split name (e.g., 'train') from the file path.
        """
        path = self.data_list[idx % len(self.data_list)]
        return os.path.basename(os.path.dirname(os.path.dirname(path)))


@DATASETS.register_module()
class ScanObjectNNDataset(DefaultDataset):
    """
    ScanObjectNN Dataset.
    Loads data from HDF5 files into memory.
    """

    def __init__(
        self,
        class_names,
        if_color=True,
        if_normal=True,
        split="train",
        data_root="data/dataset",
        transform=None,
        test_mode=False,
        test_cfg=None,
        cache=False,
        ignore_index=-1,
        loop=1,
    ):
        self.if_color = if_color
        self.if_normal = if_normal
        self.class_names = class_names
        self.data_root = data_root
        self.split = split
        self.transform = Compose(transform)
        self.cache = cache
        self.ignore_index = ignore_index
        self.loop = (
            loop if not test_mode else 1
        )  # force make loop = 1 while in test mode
        self.test_mode = test_mode
        self.test_cfg = test_cfg if test_mode else None

        if test_mode:
            self.post_transform = Compose(self.test_cfg.post_transform)
            self.aug_transform = [Compose(aug) for aug in self.test_cfg.aug_transform]

        self.data_list = self.get_data_list()
        logger = get_root_logger()
        logger.info(
            "Totally {} x {} samples in {} {} set.".format(
                len(self.data_list), self.loop, os.path.basename(self.data_root), split
            )
        )

    @staticmethod
    def get_normals(center, coords):
        Cs = np.repeat(center.reshape((1, -1)), coords.shape[0], axis=0)
        view_dirs = coords - Cs
        view_dirs = view_dirs / (
            np.linalg.norm(view_dirs, axis=-1, keepdims=True) + 1e-6
        )
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(coords)
        pcd.estimate_normals()
        normals = np.asarray(pcd.normals)
        dot_product = np.sum(normals * view_dirs, axis=-1)
        flip_mask = dot_product < 0
        normals[flip_mask] = -normals[flip_mask]
        normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True)
        return normals

    def get_data_list(self):
        """
        Called by the parent constructor.
        Here, we load the HDF5 file into memory and return a list of sample indices.
        """
        if isinstance(self.split, str):
            split_list = [self.split]
        elif isinstance(self.split, Sequence):
            split_list = self.split
        else:
            raise NotImplementedError

        self.points = []
        self.labels = []
        for split in split_list:
            if split == "train":
                h5_path = os.path.join(self.data_root, "training_objectdataset.h5")
            elif split == "test":
                h5_path = os.path.join(self.data_root, "test_objectdataset.h5")
            else:
                raise NotImplementedError(f"Split {self.split} not supported.")

            with h5py.File(h5_path, "r") as h5:
                self.points.append(np.array(h5["data"]).astype(np.float32))
                # Use np.int64 for labels, as it's the default for PyTorch's CrossEntropyLoss
                self.labels.append(np.array(h5["label"]).astype(np.int64))
        self.points = np.concatenate(self.points, axis=0)
        self.labels = np.concatenate(self.labels, axis=0)
        # The data_list is now a list of indices, from 0 to num_samples-1.
        return list(range(self.points.shape[0]))

    def get_data(self, idx):
        """
        Overrides the parent method.
        Instead of loading from a file path, we get the data from the pre-loaded numpy arrays.
        """
        data_idx = self.data_list[idx % len(self.data_list)]
        coord = self.points[data_idx].copy()
        if self.if_color:
            color = np.zeros_like(coord)
        if self.if_normal:
            normal = self.get_normals(np.mean(coord, axis=0), coord)
        label = self.labels[data_idx]
        category = np.array([label])

        data_dict = {
            "coord": coord,
            "color": color,
            "normal": normal,
            "category": category,
            "name": self.get_data_name(idx),
        }

        return data_dict

    def get_data_name(self, idx):
        """
        Overrides the parent method.
        Returns a unique name for the sample at the given index.
        """
        real_idx = idx % len(self.data_list)
        return f"{self.split}-{real_idx}"

    def get_split_name(self, idx):
        return self.split

    def prepare_test_data(self, idx):
        assert idx < len(self.data_list)
        data_dict = self.get_data(idx)
        category = data_dict.pop("category")
        data_dict = self.transform(data_dict)
        data_dict_list = []
        for aug in self.aug_transform:
            data_dict_list.append(aug(copy.deepcopy(data_dict)))
        for i in range(len(data_dict_list)):
            data_dict_list[i] = self.post_transform(data_dict_list[i])
        data_dict = dict(
            voting_list=data_dict_list,
            category=category,
            name=self.get_data_name(idx),
        )
        return data_dict


@DATASETS.register_module()
class ScanObjectNNHardestDataset(ScanObjectNNDataset):
    """
    ScanObjectNN "Hardest" variant.
    Inherits from our refactored ScanObjectNN and only overrides the file loading logic.
    """

    def get_data_list(self):
        """
        Overrides the parent method to load from the augmented HDF5 files.
        """
        if self.split == "train":
            h5_path = os.path.join(
                self.data_root, "training_objectdataset_augmentedrot_scale75.h5"
            )
        elif self.split == "test":
            h5_path = os.path.join(
                self.data_root, "test_objectdataset_augmentedrot_scale75.h5"
            )
        else:
            raise NotImplementedError(f"Split {self.split} not supported.")

        with h5py.File(h5_path, "r") as h5:
            self.points = np.array(h5["data"]).astype(np.float32)
            self.labels = np.array(h5["label"]).astype(np.int64)

        return list(range(self.points.shape[0]))