data.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: Yue Wang
@Contact: yuewangx@mit.edu
@File: data.py
@Time: 2018/10/13 6:21 PM

Modified by 
@Author: An Tao, Pengliang Ji, Ziyi Wu
@Contact: ta19@mails.tsinghua.edu.cn, jpl1723@buaa.edu.cn, dazitu616@gmail.com
@Time: 2022/7/30 7:49 PM
"""


import os
import sys
import glob
import h5py
import numpy as np
import torch
import json
import cv2
import pickle
from torch.utils.data import Dataset


def download_modelnet40():
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR = os.path.join(BASE_DIR, 'data')
    if not os.path.exists(DATA_DIR):
        os.mkdir(DATA_DIR)
    if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')):
        url = 'https://huggingface.co/datasets/Msun/modelnet40/resolve/d5dc795541800feeb7a4b3bd3142729a0d2adf7a/modelnet40_ply_hdf5_2048.zip?download=true'
        zipfile = os.path.basename(url.split('?')[0])  # Extract filename from URL
        os.system(f'wget --no-check-certificate "{url}" -O {zipfile}')
        os.system(f'unzip {zipfile}')
        os.system(f'mv modelnet40_ply_hdf5_2048 {DATA_DIR}')
        os.system(f'rm {zipfile}')


def download_shapenetpart():
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR = os.path.join(BASE_DIR, 'data')
    if not os.path.exists(DATA_DIR):
        os.mkdir(DATA_DIR)
    if not os.path.exists(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data')):
        www = 'https://shapenet.cs.stanford.edu/media/shapenet_part_seg_hdf5_data.zip'
        zipfile = os.path.basename(www)
        os.system('wget --no-check-certificate %s; unzip %s' % (www, zipfile))
        os.system('mv %s %s' % ('hdf5_data', os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data')))
        os.system('rm %s' % (zipfile))


# https://cvg-data.inf.ethz.ch/s3dis/
# https://cvg-data.inf.ethz.ch/s3dis/Stanford3dDataset_v1.2.mat
# https://cvg-data.inf.ethz.ch/s3dis/Stanford3dDataset_v1.2_Aligned_Version.zip
def download_S3DIS():
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR = os.path.join(BASE_DIR, 'data')
    if not os.path.exists(DATA_DIR):
        os.mkdir(DATA_DIR)
    if not os.path.exists(os.path.join(DATA_DIR, 'indoor3d_sem_seg_hdf5_data')):
        www = 'https://shapenet.cs.stanford.edu/media/indoor3d_sem_seg_hdf5_data.zip'
        zipfile = os.path.basename(www)
        os.system('wget %s --no-check-certificate; unzip %s' % (www, zipfile))
        os.system('mv %s %s' % ('indoor3d_sem_seg_hdf5_data', DATA_DIR))
        os.system('rm %s' % (zipfile))
    if not os.path.exists(os.path.join(DATA_DIR, 'Stanford3dDataset_v1.2_Aligned_Version')):
        if not os.path.exists(os.path.join(DATA_DIR, 'Stanford3dDataset_v1.2_Aligned_Version.zip')):
            print('Please download Stanford3dDataset_v1.2_Aligned_Version.zip \
                from https://goo.gl/forms/4SoGp4KtH1jfRqEj2 and place it under data/')
            sys.exit(0)
        else:
            zippath = os.path.join(DATA_DIR, 'Stanford3dDataset_v1.2_Aligned_Version.zip')
            os.system('unzip %s' % (zippath))
            os.system('mv %s %s' % ('Stanford3dDataset_v1.2_Aligned_Version', DATA_DIR))
            os.system('rm %s' % (zippath))


def load_data_cls(partition):
    download_modelnet40()
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR = os.path.join(BASE_DIR, 'data')
    all_data = []
    all_label = []
    for h5_name in glob.glob(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', '*%s*.h5'%partition)):
        f = h5py.File(h5_name, 'r+')
        data = f['data'][:].astype('float32')
        label = f['label'][:].astype('int64')
        f.close()
        all_data.append(data)
        all_label.append(label)
    all_data = np.concatenate(all_data, axis=0)
    all_label = np.concatenate(all_label, axis=0)
    return all_data, all_label


def load_data_partseg(partition):
    download_shapenetpart()
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR = os.path.join(BASE_DIR, 'data')
    all_data = []
    all_label = []
    all_seg = []
    if partition == 'trainval':
        file = glob.glob(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data', '*train*.h5')) \
               + glob.glob(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data', '*val*.h5'))
    else:
        file = glob.glob(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data', '*%s*.h5'%partition))
    for h5_name in file:
        f = h5py.File(h5_name, 'r+')
        data = f['data'][:].astype('float32')
        label = f['label'][:].astype('int64')
        seg = f['pid'][:].astype('int64')
        f.close()
        all_data.append(data)
        all_label.append(label)
        all_seg.append(seg)
    all_data = np.concatenate(all_data, axis=0)
    all_label = np.concatenate(all_label, axis=0)
    all_seg = np.concatenate(all_seg, axis=0)
    return all_data, all_label, all_seg


def prepare_test_data_semseg():
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR = os.path.join(BASE_DIR, 'data')
    if not os.path.exists(os.path.join(DATA_DIR, 'stanford_indoor3d')):
        os.system('python prepare_data/collect_indoor3d_data.py')
    if not os.path.exists(os.path.join(DATA_DIR, 'indoor3d_sem_seg_hdf5_data_test')):
        os.system('python prepare_data/gen_indoor3d_h5.py')


def load_data_semseg(partition, test_area):
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR = os.path.join(BASE_DIR, 'data')
    download_S3DIS()
    prepare_test_data_semseg()
    if partition == 'train':
        data_dir = os.path.join(DATA_DIR, 'indoor3d_sem_seg_hdf5_data')
    else:
        data_dir = os.path.join(DATA_DIR, 'indoor3d_sem_seg_hdf5_data_test')
    with open(os.path.join(data_dir, "all_files.txt")) as f:
        all_files = [line.rstrip() for line in f]
    with open(os.path.join(data_dir, "room_filelist.txt")) as f:
        room_filelist = [line.rstrip() for line in f]
    data_batchlist, label_batchlist = [], []
    for f in all_files:
        file = h5py.File(os.path.join(DATA_DIR, f), 'r+')
        data = file["data"][:]
        label = file["label"][:]
        data_batchlist.append(data)
        label_batchlist.append(label)
    data_batches = np.concatenate(data_batchlist, 0)
    seg_batches = np.concatenate(label_batchlist, 0)
    test_area_name = "Area_" + test_area
    train_idxs, test_idxs = [], []
    for i, room_name in enumerate(room_filelist):
        if test_area_name in room_name:
            test_idxs.append(i)
        else:
            train_idxs.append(i)
    if partition == 'train':
        all_data = data_batches[train_idxs, ...]
        all_seg = seg_batches[train_idxs, ...]
    else:
        all_data = data_batches[test_idxs, ...]
        all_seg = seg_batches[test_idxs, ...]
    return all_data, all_seg


def load_color_partseg():
    colors = []
    labels = []
    f = open("prepare_data/meta/partseg_colors.txt")
    for line in json.load(f):
        colors.append(line['color'])
        labels.append(line['label'])
    partseg_colors = np.array(colors)
    partseg_colors = partseg_colors[:, [2, 1, 0]]
    partseg_labels = np.array(labels)
    font = cv2.FONT_HERSHEY_SIMPLEX
    img_size = 1350
    img = np.zeros((1350, 1890, 3), dtype="uint8")
    cv2.rectangle(img, (0, 0), (1900, 1900), [255, 255, 255], thickness=-1)
    column_numbers = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3]
    column_gaps = [320, 320, 300, 300, 285, 285]
    color_size = 64
    color_index = 0
    label_index = 0
    row_index = 16
    for row in range(0, img_size):
        column_index = 32
        for column in range(0, img_size):
            color = partseg_colors[color_index]
            label = partseg_labels[label_index]
            length = len(str(label))
            cv2.rectangle(img, (column_index, row_index), (column_index + color_size, row_index + color_size),
                          color=(int(color[0]), int(color[1]), int(color[2])), thickness=-1)
            img = cv2.putText(img, label, (column_index + int(color_size * 1.15), row_index + int(color_size / 2)),
                              font,
                              0.76, (0, 0, 0), 2)
            column_index = column_index + column_gaps[column]
            color_index = color_index + 1
            label_index = label_index + 1
            if color_index >= 50:
                cv2.imwrite("prepare_data/meta/partseg_colors.png", img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
                return np.array(colors)
            elif (column + 1 >= column_numbers[row]):
                break
        row_index = row_index + int(color_size * 1.3)
        if (row_index >= img_size):
            break


def load_color_semseg():
    colors = []
    labels = []
    f = open("prepare_data/meta/semseg_colors.txt")
    for line in json.load(f):
        colors.append(line['color'])
        labels.append(line['label'])
    semseg_colors = np.array(colors)
    semseg_colors = semseg_colors[:, [2, 1, 0]]
    partseg_labels = np.array(labels)
    font = cv2.FONT_HERSHEY_SIMPLEX
    img_size = 1500
    img = np.zeros((500, img_size, 3), dtype="uint8")
    cv2.rectangle(img, (0, 0), (img_size, 750), [255, 255, 255], thickness=-1)
    color_size = 64
    color_index = 0
    label_index = 0
    row_index = 16
    for _ in range(0, img_size):
        column_index = 32
        for _ in range(0, img_size):
            color = semseg_colors[color_index]
            label = partseg_labels[label_index]
            length = len(str(label))
            cv2.rectangle(img, (column_index, row_index), (column_index + color_size, row_index + color_size),
                          color=(int(color[0]), int(color[1]), int(color[2])), thickness=-1)
            img = cv2.putText(img, label, (column_index + int(color_size * 1.15), row_index + int(color_size / 2)),
                              font,
                              0.7, (0, 0, 0), 2)
            column_index = column_index + 200
            color_index = color_index + 1
            label_index = label_index + 1
            if color_index >= 13:
                cv2.imwrite("prepare_data/meta/semseg_colors.png", img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
                return np.array(colors)
            elif (column_index >= 1280):
                break
        row_index = row_index + int(color_size * 1.3)
        if (row_index >= img_size):
            break  
    

def translate_pointcloud(pointcloud):
    xyz1 = np.random.uniform(low=2./3., high=3./2., size=[3])
    xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3])
       
    translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32')
    return translated_pointcloud


def jitter_pointcloud(pointcloud, sigma=0.01, clip=0.02):
    N, C = pointcloud.shape
    pointcloud += np.clip(sigma * np.random.randn(N, C), -1*clip, clip)
    return pointcloud


def rotate_pointcloud(pointcloud):
    theta = np.pi*2 * np.random.uniform()
    rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
    pointcloud[:,[0,2]] = pointcloud[:,[0,2]].dot(rotation_matrix) # random rotation (x,z)
    return pointcloud


class ModelNet40(Dataset):
    def __init__(self, num_points, partition='train'):
        self.data, self.label = load_data_cls(partition)
        self.num_points = num_points
        self.partition = partition        

    def __getitem__(self, item):
        pointcloud = self.data[item][:self.num_points]
        label = self.label[item]
        if self.partition == 'train':
            pointcloud = translate_pointcloud(pointcloud)
            np.random.shuffle(pointcloud)
        return pointcloud, label

    def __len__(self):
        return self.data.shape[0]


class ShapeNetPart(Dataset):
    def __init__(self, num_points, partition='train', class_choice=None):
        self.data, self.label, self.seg = load_data_partseg(partition)
        self.cat2id = {'airplane': 0, 'bag': 1, 'cap': 2, 'car': 3, 'chair': 4, 
                       'earphone': 5, 'guitar': 6, 'knife': 7, 'lamp': 8, 'laptop': 9, 
                       'motor': 10, 'mug': 11, 'pistol': 12, 'rocket': 13, 'skateboard': 14, 'table': 15}
        self.seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3]
        self.index_start = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47]
        self.num_points = num_points
        self.partition = partition        
        self.class_choice = class_choice
        self.partseg_colors = load_color_partseg()
        
        if self.class_choice != None:
            id_choice = self.cat2id[self.class_choice]
            indices = (self.label == id_choice).squeeze()
            self.data = self.data[indices]
            self.label = self.label[indices]
            self.seg = self.seg[indices]
            self.seg_num_all = self.seg_num[id_choice]
            self.seg_start_index = self.index_start[id_choice]
        else:
            self.seg_num_all = 50
            self.seg_start_index = 0
            
      
    def __getitem__(self, item):
        pointcloud = self.data[item][:self.num_points]
        label = self.label[item]
        seg = self.seg[item][:self.num_points]
        if self.partition == 'trainval':
            # pointcloud = translate_pointcloud(pointcloud)
            indices = list(range(pointcloud.shape[0]))
            np.random.shuffle(indices)
            pointcloud = pointcloud[indices]
            seg = seg[indices]
        return pointcloud, label, seg

    def __len__(self):
        return self.data.shape[0]


class S3DIS(Dataset):
    def __init__(self, num_points=4096, partition='train', test_area='1'):
        self.data, self.seg = load_data_semseg(partition, test_area)
        self.num_points = num_points
        self.partition = partition    
        self.semseg_colors = load_color_semseg()

    def __getitem__(self, item):
        pointcloud = self.data[item][:self.num_points]
        seg = self.seg[item][:self.num_points]
        if self.partition == 'train':
            indices = list(range(pointcloud.shape[0]))
            np.random.shuffle(indices)
            pointcloud = pointcloud[indices]
            seg = seg[indices]
        seg = torch.LongTensor(seg)
        return pointcloud, seg

    def __len__(self):
        return self.data.shape[0]


    


# Specs:
#   pointcloud.shape == (num_points, 3)
#   seg.shape == (num_points,)
class TractorsAndCombines(Dataset):
    def __init__(self, root, split='train', augment=False):
        self.root = root
        self.split = split
        self.augment = augment # Enable/disable data augmentation
        self.sem_classes = {'Ground': [0], 'Tractor': [1], 'Combine Harvester': [2]}
        self.processed_dir = os.path.join(root, 'processed_dgcnn')
        
        if os.path.exists(self.processed_dir):
            print('Processed data found!')
        else:
            raise RuntimeError('Processed data not found! Please run "prepare_tractors.py" first!')
        
        # Get the pointcloud and label file paths and sort them
        if split == 'train' or split == 'validate' or split == 'test':
            self.pointcloud_files = sorted([os.path.join(self.processed_dir, split, 'points', f) 
                                            for f in os.listdir(os.path.join(self.processed_dir, split, 'points'))])
            self.label_files = sorted([os.path.join(self.processed_dir, split, 'labels', f)
                                       for f in os.listdir(os.path.join(self.processed_dir, split, 'labels'))])
        else:
            print(f'Invalid split: "{self.split}"! Please use "train", "validate", or "test"!')
            sys.exit(1)
        
    # Set some class properties
    @property
    def classes(self):
        return self.sem_classes
    
    @property
    def num_classes(self):
        return len(self.sem_classes)
    
    @property
    def num_points(self):
        return np.load(self.pointcloud_files[0]).shape[0]
    
    @property
    def num_features(self):
        return np.load(self.pointcloud_files[0]).shape[1]
    
    
    def __getitem__(self, index):
        # Load the pointcloud
        pointcloud: np.ndarray = np.load(self.pointcloud_files[index]).astype(np.float32) # Should be saved as np.float16
        
        # Load the label
        seg: np.ndarray = np.load(self.label_files[index]) # Should be saved as np.int8
        
        # Apply data augmentation for the 'train' split only
        if self.augment and self.split == 'train':
            pointcloud, seg = self.augment_pointcloud(pointcloud, seg)
        
        pointcloud = torch.tensor(pointcloud, dtype=torch.float32)
        seg = torch.tensor(seg, dtype=torch.int64)
        seg = torch.squeeze(input=seg) # Ensure shape is [num_points] instead of [num_points, 1]
        
        return pointcloud, seg
    
    def __len__(self):
        return len(self.pointcloud_files)
    
    def augment_pointcloud(self, pointcloud, seg):
        """Apply augmentations to the point cloud and corresponding labels."""
        # Random rotation around z-axis
        theta = np.random.uniform(0, 2 * np.pi)
        rotation_matrix = np.array([
            [np.cos(theta), -np.sin(theta), 0],
            [np.sin(theta), np.cos(theta), 0],
            [0, 0, 1]
        ])
        pointcloud[:, :3] = pointcloud[:, :3] @ rotation_matrix.T

        # Add Gaussian noise
        noise = np.random.normal(0, 0.01, size=pointcloud.shape)
        pointcloud += noise

        # Random scaling
        scale = np.random.uniform(0.8, 1.2)
        pointcloud[:, :3] *= scale

        # Random translation
        translation = np.random.uniform(-0.1, 0.1, size=(1, 3))
        pointcloud[:, :3] += translation

        return pointcloud, seg



class ScanNet(Dataset):
    def __init__(self, num_point=8192, partition='train', 
                 data_root='scannet', classes=20, block_size=1.5,
                 sample_rate=1.0, transform=None, use_rgb=False):
        self.partition = partition
        self.num_point = num_point
        self.block_size = block_size
        self.transform = transform
        xyz_all = []
        label_all = []
        if not isinstance(partition, list):
            partition = [partition]
        for i in partition:
            data_file = os.path.join(
                    data_root, 'scannet_{}_rgb21c_pointid.pickle'.format(i))
            file_pickle = open(data_file, 'rb')
            _xyz_all = pickle.load(file_pickle)
            _label_all = pickle.load(file_pickle)
            file_pickle.close()
            xyz_all.append(_xyz_all)
            label_all.append(_label_all)
        xyz_all = np.hstack(xyz_all)
        label_all = np.hstack(label_all)
        self.label_all = []  # for change 0-20 to 0-19 + 255
        self.room_coord_min, self.room_coord_max = [], []
        num_point_all = []
        for index in range(len(xyz_all)):
            xyz, label = xyz_all[index], label_all[index]  # xyzrgb, N*6; l, N
            coord_min, coord_max = np.amin(
                xyz, axis=0)[:3], np.amax(xyz, axis=0)[:3]
            self.room_coord_min.append(coord_min)
            self.room_coord_max.append(coord_max)
            num_point_all.append(label.size)
            # we have set all ignore_class to 0
            # class 0 is also ignored
            # so we set all them as 255
            label_new = label - 1
            label_new[label == 0] = 255
            self.label_all.append(label_new.astype(np.uint8))
        sample_prob = num_point_all / np.sum(num_point_all)
        num_iter = int(np.sum(num_point_all) * sample_rate / num_point)
        room_idxs = []
        for index in range(len(xyz_all)):
            room_idxs.extend(
                [index] * int(round(sample_prob[index] * num_iter)))
        self.room_idxs = np.array(room_idxs)
        self.xyz_all = xyz_all

        # whether load RGB information
        self.use_rgb = use_rgb

        print("Totally {} samples in {} set.".format(len(self.room_idxs), partition))

    def __getitem__(self, idx):
        room_idx = self.room_idxs[idx]
        points = self.xyz_all[room_idx]  # N * 6
        if not self.use_rgb:
            points = points[:, :3]
        labels = self.label_all[room_idx]  # N
        N_points = points.shape[0]

        for i in range(10):
            center = points[np.random.choice(N_points)][:3]
            block_min = center - [self.block_size /
                                  2.0, self.block_size / 2.0, 0]
            block_max = center + [self.block_size /
                                  2.0, self.block_size / 2.0, 0]
            block_min[2] = self.room_coord_min[room_idx][2]
            block_max[2] = self.room_coord_max[room_idx][2]
            point_idxs = np.where((points[:, 0] >= block_min[0]) &
                                  (points[:, 0] <= block_max[0]) &
                                  (points[:, 1] >= block_min[1]) &
                                  (points[:, 1] <= block_max[1]))[0]
            if point_idxs.size == 0:
                continue
            vidx = np.ceil((points[point_idxs, :3] - block_min) /
                           (block_max - block_min) * [31.0, 31.0, 62.0])
            vidx = np.unique(vidx[:, 0] * 31.0 * 62.0 +
                             vidx[:, 1] * 62.0 + vidx[:, 2])
            if ((labels[point_idxs] != 255).sum() / point_idxs.size >= 0.7) and (vidx.size/31.0/31.0/62.0 >= 0.02):
                break

        if point_idxs.size >= self.num_point:
            selected_point_idxs = np.random.choice(
                point_idxs, self.num_point, replace=False)
        else:
            selected_point_idxs = np.random.choice(
                point_idxs, self.num_point, replace=True)
        # normalize
        selected_points = points[selected_point_idxs, :]  # num_point * 3/6
        num_feats = 9 if self.use_rgb else 6
        current_points = np.zeros(
            (self.num_point, num_feats))  # num_point * 6/9
        current_points[:, -3] = selected_points[:, 0] / \
            self.room_coord_max[room_idx][0]
        current_points[:, -2] = selected_points[:, 1] / \
            self.room_coord_max[room_idx][1]
        current_points[:, -1] = selected_points[:, 2] / \
            self.room_coord_max[room_idx][2]
        selected_points[:, 0] = selected_points[:, 0] - center[0]
        selected_points[:, 1] = selected_points[:, 1] - center[1]
        current_points[:, 0:3] = selected_points[:, 0:3]
        if self.use_rgb:
            current_points[:, 3:6] = selected_points[:, 3:6] / 255.
        current_labels = labels[selected_point_idxs]
        if self.transform is not None:
            current_points, current_labels = self.transform(
                current_points, current_labels)
        return current_points, current_labels

    def __len__(self):
        return len(self.room_idxs)

if __name__ == '__main__':
    train = ModelNet40(1024)
    test = ModelNet40(1024, 'test')
    data, label = train[0]
    print(data.shape)
    print(label.shape)

    trainval = ShapeNetPart(2048, 'trainval')
    test = ShapeNetPart(2048, 'test')
    data, label, seg = trainval[0]
    print(data.shape)
    print(label.shape)
    print(seg.shape)

    train = S3DIS(4096)
    test = S3DIS(4096, 'test')
    data, seg = train[0]
    print(data.shape)
    print(seg.shape)

    train = ScanNet(8192)
    test = ScanNet(8192, 'test')
    data, seg = train[0]
    print(data.shape)
    print(seg.shape)