🐛 missing imports

This commit is contained in:
Laurent FAINSIN 2023-05-16 11:04:04 +02:00
parent d0cdb8e4ee
commit 6e682482d2

View file

@ -30,9 +30,9 @@ import torch
import warnings import warnings
from multiprocessing import Lock from multiprocessing import Lock
from utils.ply import write_ply, read_ply
# OS functions # OS functions
from os import listdir from os import listdir, makedirs
from os.path import exists, join, isdir from os.path import exists, join, isdir
# Dataset parent class # Dataset parent class
@ -757,6 +757,7 @@ class S3DISDataset(PointCloudDataset):
cloud_classes = np.vstack((cloud_classes, object_classes)) cloud_classes = np.vstack((cloud_classes, object_classes))
# Save as ply # Save as ply
write_ply( write_ply(
cloud_file, cloud_file,
(cloud_points, cloud_colors, cloud_classes), (cloud_points, cloud_colors, cloud_classes),
@ -830,6 +831,8 @@ class S3DISDataset(PointCloudDataset):
sub_labels = np.squeeze(sub_labels) sub_labels = np.squeeze(sub_labels)
# Get chosen neighborhoods # Get chosen neighborhoods
from sklearn.neighbors import KDTree
search_tree = KDTree(sub_points, leaf_size=10) search_tree = KDTree(sub_points, leaf_size=10)
# search_tree = nnfln.KDTree(n_neighbors=1, metric='L2', leaf_size=10) # search_tree = nnfln.KDTree(n_neighbors=1, metric='L2', leaf_size=10)
# search_tree.fit(sub_points) # search_tree.fit(sub_points)