<filename>pointNet/framework/dataset.py
import torch
import os
import json
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import numpy as np
from tqdm import tqdm


def read_pcd_from_file(file):
    np_pts = np.zeros(0)
    with open(file, 'r') as f:
        pts = []
        for line in f:
            one_pt = list(map(float, line[:-1].split(',')))
            pts.append(one_pt[:3])
        np_pts = np.array(pts)
    return np_pts


def read_file_names_from_file(file):
  with open(file, 'r') as f:
    files = []
    for line in f:
      files.append(line.split('\n')[0])
  return files


class PointNetDataset(Dataset):
  def __init__(self, root_dir, train):
    super(PointNetDataset, self).__init__()

    self._train = train
    self._classes = []

    self._features = []
    self._labels = []

    self.load(root_dir)

  def classes(self):
    return self._classes

  def __len__(self):
    return len(self._features)
  
  def __getitem__(self, idx):
    # import ipdb;ipdb.set_trace()
    feature, label = self._features[idx], self._labels[idx]
    
    # normalize feature
    feature = feature - np.mean(feature,axis=0)
    scale = np.max(np.sqrt(np.sum(feature ** 2, axis = 1)),0)
    feature = feature / scale
    
    # data augmentation, random rotation
    theta = np.random.uniform(0, np.pi * 2)
    rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
    feature[:,[0,2]] = feature[:,[0,2]].dot(rotation_matrix)


    # data augmentation, random jitter
    feature += np.random.normal(0, 0.02, size=feature.shape)
    feature = torch.Tensor(feature.T)

    l_lable = [0 for _ in range(len(self._classes))]
    l_lable[self._classes.index(label)] = 1
    label = torch.Tensor(l_lable)

    return feature, label
  
  def load(self, root_dir):
    # import ipdb;ipdb.set_trace()
    things = os.listdir(root_dir)
    files = []
    for f in things:
      if self._train == 0:
        if f == 'modelnet40_train.txt':
          files = read_file_names_from_file(root_dir + '/' + f)
      elif self._train == 1:
        if f == 'modelnet40_test.txt':
          files = read_file_names_from_file(root_dir + '/' + f)
      if f == "modelnet40_shape_names.txt":
        self._classes = read_file_names_from_file(root_dir + '/' + f)
    tmp_classes = []
    # import ipdb;ipdb.set_trace()
    for file in tqdm(files):
      num = file.split("_")[-1]
      kind = file.split("_" + num)[0]
      if kind not in tmp_classes:
        tmp_classes.append(kind)
      pcd_file = root_dir + '/' + kind + '/' + file + '.txt'
      np_pts = read_pcd_from_file(pcd_file)
      # print(np_pts.shape) # (10000, 3)
      self._features.append(np_pts)
      self._labels.append(kind)
    if self._train == 0:
      print("There are " + str(len(self._labels)) + " train files.")
    elif self._train == 1:
      print("There are " + str(len(self._labels)) + " test files.")

def feature_transform_regularizer(trans):
    d = trans.size()[1]
    batchsize = trans.size()[0]
    I = torch.eye(d)[None, :, :]
    if trans.is_cuda:
        I = I.cuda()
    loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2,1)) - I, dim=(1,2)))
    return loss

if __name__ == "__main__":
  train_data = PointNetDataset("../../modelnet40_normal_resampled", train=0)
  # train_data = PointNetDataset("/home/swarm/developments/point_cloud/modelnet40_normal_resampled/",train=0)
  # train_data = PointNetDataset("../../modelnet40_normal_resampled",train=0)
  train_loader = DataLoader(train_data, batch_size=2, shuffle=True)
  # import ipdb;ipdb.set_trace()
  cnt = 0
  for pts, label in train_loader:
    # print(pts.shape)
    import ipdb;ipdb.set_trace()
    print(label.shape)
    cnt += 1
    if cnt > 3:
        break
    print(train_data[0])

<filename>pointNet/framework/dataset.py
import torch
import os
import json
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import numpy as np
from tqdm import tqdm


def read_pcd_from_file(file):
    np_pts = np.zeros(0)
    with open(file, 'r') as f:
        pts = []
        for line in f:
            one_pt = list(map(float, line[:-1].split(',')))
            pts.append(one_pt[:3])
        np_pts = np.array(pts)
    return np_pts


def read_file_names_from_file(file):
  with open(file, 'r') as f:
    files = []
    for line in f:
      files.append(line.split('\n')[0])
  return files


class PointNetDataset(Dataset):
  def __init__(self, root_dir, train):
    super(PointNetDataset, self).__init__()

    self._train = train
    self._classes = []

    self._features = []
    self._labels = []

    self.load(root_dir)

  def classes(self):
    return self._classes

  def __len__(self):
    return len(self._features)
  
  def __getitem__(self, idx):
    # import ipdb;ipdb.set_trace()
    feature, label = self._features[idx], self._labels[idx]
    
    # normalize feature
    feature = feature - np.mean(feature,axis=0)
    scale = np.max(np.sqrt(np.sum(feature ** 2, axis = 1)),0)
    feature = feature / scale
    
    # data augmentation, random rotation
    theta = np.random.uniform(0, np.pi * 2)
    rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
    feature[:,[0,2]] = feature[:,[0,2]].dot(rotation_matrix)


    # data augmentation, random jitter
    feature += np.random.normal(0, 0.02, size=feature.shape)
    feature = torch.Tensor(feature.T)

    l_lable = [0 for _ in range(len(self._classes))]
    l_lable[self._classes.index(label)] = 1
    label = torch.Tensor(l_lable)

    return feature, label
  
  def load(self, root_dir):
    # import ipdb;ipdb.set_trace()
    things = os.listdir(root_dir)
    files = []
    for f in things:
      if self._train == 0:
        if f == 'modelnet40_train.txt':
          files = read_file_names_from_file(root_dir + '/' + f)
      elif self._train == 1:
        if f == 'modelnet40_test.txt':
          files = read_file_names_from_file(root_dir + '/' + f)
      if f == "modelnet40_shape_names.txt":
        self._classes = read_file_names_from_file(root_dir + '/' + f)
    tmp_classes = []
    # import ipdb;ipdb.set_trace()
    for file in tqdm(files):
      num = file.split("_")[-1]
      kind = file.split("_" + num)[0]
      if kind not in tmp_classes:
        tmp_classes.append(kind)
      pcd_file = root_dir + '/' + kind + '/' + file + '.txt'
      np_pts = read_pcd_from_file(pcd_file)
      # print(np_pts.shape) # (10000, 3)
      self._features.append(np_pts)
      self._labels.append(kind)
    if self._train == 0:
      print("There are " + str(len(self._labels)) + " train files.")
    elif self._train == 1:
      print("There are " + str(len(self._labels)) + " test files.")

def feature_transform_regularizer(trans):
    d = trans.size()[1]
    batchsize = trans.size()[0]
    I = torch.eye(d)[None, :, :]
    if trans.is_cuda:
        I = I.cuda()
    loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2,1)) - I, dim=(1,2)))
    return loss

if __name__ == "__main__":
  train_data = PointNetDataset("../../modelnet40_normal_resampled", train=0)
  # train_data = PointNetDataset("/home/swarm/developments/point_cloud/modelnet40_normal_resampled/",train=0)
  # train_data = PointNetDataset("../../modelnet40_normal_resampled",train=0)
  train_loader = DataLoader(train_data, batch_size=2, shuffle=True)
  # import ipdb;ipdb.set_trace()
  cnt = 0
  for pts, label in train_loader:
    # print(pts.shape)
    import ipdb;ipdb.set_trace()
    print(label.shape)
    cnt += 1
    if cnt > 3:
        break
    print(train_data[0])

# import ipdb;ipdb.set_trace() # normalize feature # data augmentation, random rotation # data augmentation, random jitter # import ipdb;ipdb.set_trace() # import ipdb;ipdb.set_trace() # print(np_pts.shape) # (10000, 3) # train_data = PointNetDataset("/home/swarm/developments/point_cloud/modelnet40_normal_resampled/",train=0) # train_data = PointNetDataset("../../modelnet40_normal_resampled",train=0) # import ipdb;ipdb.set_trace() # print(pts.shape)