"tests/git@developer.sourcefind.cn:OpenDAS/d2go.git" did not exist on "27f2b21d89fedd6c8b886066a3c27123c607b9f5"
Commit 9671da9d authored by liyinhao's avatar liyinhao
Browse files

change class names

parent c8fa96f8
...@@ -2,7 +2,7 @@ import os ...@@ -2,7 +2,7 @@ import os
import pickle import pickle
from pathlib import Path from pathlib import Path
from tools.data_converter.scannet_data_utils import ScannetObject from tools.data_converter.scannet_data_utils import ScanNetData
def create_scannet_info_file(data_path, pkl_prefix='scannet', save_path=None): def create_scannet_info_file(data_path, pkl_prefix='scannet', save_path=None):
...@@ -14,8 +14,8 @@ def create_scannet_info_file(data_path, pkl_prefix='scannet', save_path=None): ...@@ -14,8 +14,8 @@ def create_scannet_info_file(data_path, pkl_prefix='scannet', save_path=None):
assert os.path.exists(save_path) assert os.path.exists(save_path)
train_filename = save_path / f'{pkl_prefix}_infos_train.pkl' train_filename = save_path / f'{pkl_prefix}_infos_train.pkl'
val_filename = save_path / f'{pkl_prefix}_infos_val.pkl' val_filename = save_path / f'{pkl_prefix}_infos_val.pkl'
train_dataset = ScannetObject(root_path=data_path, split='train') train_dataset = ScanNetData(root_path=data_path, split='train')
val_dataset = ScannetObject(root_path=data_path, split='val') val_dataset = ScanNetData(root_path=data_path, split='val')
scannet_infos_train = train_dataset.get_scannet_infos(has_label=True) scannet_infos_train = train_dataset.get_scannet_infos(has_label=True)
with open(train_filename, 'wb') as f: with open(train_filename, 'wb') as f:
pickle.dump(scannet_infos_train, f) pickle.dump(scannet_infos_train, f)
......
...@@ -3,7 +3,7 @@ import os ...@@ -3,7 +3,7 @@ import os
import numpy as np import numpy as np
class ScannetObject(object): class ScanNetData(object):
''' Load and parse object data ''' ''' Load and parse object data '''
def __init__(self, root_path, split='train'): def __init__(self, root_path, split='train'):
...@@ -79,7 +79,4 @@ class ScannetObject(object): ...@@ -79,7 +79,4 @@ class ScannetObject(object):
else self.sample_id_list else self.sample_id_list
with futures.ThreadPoolExecutor(num_workers) as executor: with futures.ThreadPoolExecutor(num_workers) as executor:
infos = executor.map(process_single_scene, sample_id_list) infos = executor.map(process_single_scene, sample_id_list)
# infos = list()
# for sample in sample_id_list:
# infos.append(process_single_scene(sample))
return list(infos) return list(infos)
...@@ -2,7 +2,7 @@ import os ...@@ -2,7 +2,7 @@ import os
import pickle import pickle
from pathlib import Path from pathlib import Path
from tools.data_converter.sunrgbd_data_utils import SUNRGBDObject from tools.data_converter.sunrgbd_data_utils import SUNRGBDData
def create_sunrgbd_info_file(data_path, def create_sunrgbd_info_file(data_path,
...@@ -17,10 +17,9 @@ def create_sunrgbd_info_file(data_path, ...@@ -17,10 +17,9 @@ def create_sunrgbd_info_file(data_path,
assert os.path.exists(save_path) assert os.path.exists(save_path)
train_filename = save_path / f'{pkl_prefix}_infos_train.pkl' train_filename = save_path / f'{pkl_prefix}_infos_train.pkl'
val_filename = save_path / f'{pkl_prefix}_infos_val.pkl' val_filename = save_path / f'{pkl_prefix}_infos_val.pkl'
train_dataset = SUNRGBDObject( train_dataset = SUNRGBDData(
root_path=data_path, split='train', use_v1=use_v1) root_path=data_path, split='train', use_v1=use_v1)
val_dataset = SUNRGBDObject( val_dataset = SUNRGBDData(root_path=data_path, split='val', use_v1=use_v1)
root_path=data_path, split='val', use_v1=use_v1)
sunrgbd_infos_train = train_dataset.get_sunrgbd_infos(has_label=True) sunrgbd_infos_train = train_dataset.get_sunrgbd_infos(has_label=True)
with open(train_filename, 'wb') as f: with open(train_filename, 'wb') as f:
pickle.dump(sunrgbd_infos_train, f) pickle.dump(sunrgbd_infos_train, f)
......
...@@ -17,7 +17,7 @@ def random_sampling(pc, num_sample, replace=None, return_choices=False): ...@@ -17,7 +17,7 @@ def random_sampling(pc, num_sample, replace=None, return_choices=False):
return pc[choices] return pc[choices]
class SUNObject3d(object): class SUNRGBDInstance(object):
def __init__(self, line): def __init__(self, line):
data = line.split(' ') data = line.split(' ')
...@@ -43,7 +43,7 @@ class SUNObject3d(object): ...@@ -43,7 +43,7 @@ class SUNObject3d(object):
]) ])
class SUNRGBDObject(object): class SUNRGBDData(object):
''' Load and parse object data ''' ''' Load and parse object data '''
def __init__(self, root_path, split='train', use_v1=False): def __init__(self, root_path, split='train', use_v1=False):
...@@ -57,7 +57,7 @@ class SUNRGBDObject(object): ...@@ -57,7 +57,7 @@ class SUNRGBDObject(object):
self.cat2label = {cat: self.classes.index(cat) for cat in self.classes} self.cat2label = {cat: self.classes.index(cat) for cat in self.classes}
self.label2cat = { self.label2cat = {
label: self.classes[label] label: self.classes[label]
for label in len(self.classes) for label in range(len(self.classes))
} }
assert split in ['train', 'val', 'test'] assert split in ['train', 'val', 'test']
split_dir = os.path.join(self.root_dir, '%s_data_idx.txt' % split) split_dir = os.path.join(self.root_dir, '%s_data_idx.txt' % split)
...@@ -100,7 +100,7 @@ class SUNRGBDObject(object): ...@@ -100,7 +100,7 @@ class SUNRGBDObject(object):
def get_label_objects(self, idx): def get_label_objects(self, idx):
label_filename = os.path.join(self.label_dir, '%06d.txt' % (idx)) label_filename = os.path.join(self.label_dir, '%06d.txt' % (idx))
lines = [line.rstrip() for line in open(label_filename)] lines = [line.rstrip() for line in open(label_filename)]
objects = [SUNObject3d(line) for line in lines] objects = [SUNRGBDInstance(line) for line in lines]
return objects return objects
def get_sunrgbd_infos(self, def get_sunrgbd_infos(self,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment