forked from torch-points3d/torch-points3d
-
Notifications
You must be signed in to change notification settings - Fork 44
/
pretrained_api.py
174 lines (145 loc) · 9.35 KB
/
pretrained_api.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
import os
import logging
import urllib.request
from omegaconf import DictConfig
# Import building function for model and dataset
from torch_points3d.datasets.dataset_factory import instantiate_dataset
from torch_points3d.models.model_factory import instantiate_model
# Import BaseModel / BaseDataset for type checking
from torch_points3d.models.base_model import BaseModel
from torch_points3d.datasets.base_dataset import BaseDataset
from torch_points3d.utils.wandb_utils import Wandb
from torch_points3d.metrics.model_checkpoint import ModelCheckpoint
log = logging.getLogger(__name__)
DIR = os.path.dirname(os.path.realpath(__file__))
CHECKPOINT_DIR = os.path.join(DIR, "weights")
def download_file(url, out_file):
if not os.path.exists(out_file):
if not os.path.exists(os.path.dirname(out_file)):
os.makedirs(os.path.dirname(out_file))
urllib.request.urlretrieve(url, out_file)
else:
log.warning("WARNING: skipping download of existing file " + out_file)
class PretainedRegistry(object):
MODELS = {
"pointnet2_largemsg-s3dis-1": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/1e1p0csk/pointnet2_largemsg.pt",
"pointnet2_largemsg-s3dis-2": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/2i499g2e/pointnet2_largemsg.pt",
"pointnet2_largemsg-s3dis-3": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/1gyokj69/pointnet2_largemsg.pt",
"pointnet2_largemsg-s3dis-4": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/1ejjs4s2/pointnet2_largemsg.pt",
"pointnet2_largemsg-s3dis-5": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/etxij0j6/pointnet2_largemsg.pt",
"pointnet2_largemsg-s3dis-6": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/8n8t391d/pointnet2_largemsg.pt",
"pointgroup-scannet": "https://api.wandb.ai/files/nicolas/panoptic/2ta6vfu2/PointGroup.pt",
"minkowski-res16-s3dis-1": "https://api.wandb.ai/files/nicolas/s3dis-benchmark/1fyr7ri9/Res16UNet34C.pt",
"minkowski-res16-s3dis-2": "https://api.wandb.ai/files/nicolas/s3dis-benchmark/1gdgx2ni/Res16UNet34C.pt",
"minkowski-res16-s3dis-3": "https://api.wandb.ai/files/nicolas/s3dis-benchmark/gt3ttamp/Res16UNet34C.pt",
"minkowski-res16-s3dis-4": "https://api.wandb.ai/files/nicolas/s3dis-benchmark/36yxu3yc/Res16UNet34C.pt",
"minkowski-res16-s3dis-5": "https://api.wandb.ai/files/nicolas/s3dis-benchmark/2r0tsub1/Res16UNet34C.pt",
"minkowski-res16-s3dis-6": "https://api.wandb.ai/files/nicolas/s3dis-benchmark/30yrkk5p/Res16UNet34C.pt",
"minkowski-registration-3dmatch": "https://api.wandb.ai/files/humanpose1/registration/2wvwf92e/MinkUNet_Fragment.pt",
"minkowski-registration-kitti": "https://api.wandb.ai/files/humanpose1/KITTI/2xpy7u1i/MinkUNet_Fragment.pt",
"minkowski-registration-modelnet": "https://api.wandb.ai/files/humanpose1/modelnet/39u5v3bm/MinkUNet_Fragment.pt",
"rsconv-s3dis-1": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/2b99o12e/RSConv_MSN_S3DIS.pt",
"rsconv-s3dis-2": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/1onl4h59/RSConv_MSN_S3DIS.pt",
"rsconv-s3dis-3": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/2cau6jua/RSConv_MSN_S3DIS.pt",
"rsconv-s3dis-4": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/1qqmzgnz/RSConv_MSN_S3DIS.pt",
"rsconv-s3dis-5": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/378enxsu/RSConv_MSN_S3DIS.pt",
"rsconv-s3dis-6": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/23f4upgc/RSConv_MSN_S3DIS.pt",
"kpconv-s3dis-1": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/okiba8gp/KPConvPaper.pt",
"kpconv-s3dis-2": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/2at56wrm/KPConvPaper.pt",
"kpconv-s3dis-3": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/1ipv9lso/KPConvPaper.pt",
"kpconv-s3dis-4": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/2c13jhi0/KPConvPaper.pt",
"kpconv-s3dis-5": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/1kf8yg5s/KPConvPaper.pt",
"kpconv-s3dis-6": "https://api.wandb.ai/files/loicland/benchmark-torch-points-3d-s3dis/2ph7ejss/KPConvPaper.pt",
}
MOCK_USED_PROPERTIES = {
"pointnet2_largemsg-s3dis-1": {"feature_dimension": 4, "num_classes": 13},
"pointnet2_largemsg-s3dis-2": {"feature_dimension": 4, "num_classes": 13},
"pointnet2_largemsg-s3dis-3": {"feature_dimension": 4, "num_classes": 13},
"pointnet2_largemsg-s3dis-4": {"feature_dimension": 4, "num_classes": 13},
"pointnet2_largemsg-s3dis-5": {"feature_dimension": 4, "num_classes": 13},
"pointnet2_largemsg-s3dis-6": {"feature_dimension": 4, "num_classes": 13},
"pointgroup-scannet": {},
"rsconv-s3dis-1": {"feature_dimension": 4, "num_classes": 13},
"rsconv-s3dis-2": {"feature_dimension": 4, "num_classes": 13},
"rsconv-s3dis-3": {"feature_dimension": 4, "num_classes": 13},
"rsconv-s3dis-4": {"feature_dimension": 4, "num_classes": 13},
"rsconv-s3dis-5": {"feature_dimension": 4, "num_classes": 13},
"rsconv-s3dis-6": {"feature_dimension": 4, "num_classes": 13},
"minkowski-res16-s3dis-1": {"feature_dimension": 4, "num_classes": 13},
"minkowski-res16-s3dis-2": {"feature_dimension": 4, "num_classes": 13},
"minkowski-res16-s3dis-3": {"feature_dimension": 4, "num_classes": 13},
"minkowski-res16-s3dis-4": {"feature_dimension": 4, "num_classes": 13},
"minkowski-res16-s3dis-5": {"feature_dimension": 4, "num_classes": 13},
"minkowski-res16-s3dis-6": {"feature_dimension": 4, "num_classes": 13},
"minkowski-registration-3dmatch": {"feature_dimension": 1},
"minkowski-registration-kitti": {"feature_dimension": 1},
"minkowski-registration-modelnet": {"feature_dimension": 1},
"kpconv-s3dis-1": {"feature_dimension": 4, "num_classes": 13},
"kpconv-s3dis-2": {"feature_dimension": 4, "num_classes": 13},
"kpconv-s3dis-3": {"feature_dimension": 4, "num_classes": 13},
"kpconv-s3dis-4": {"feature_dimension": 4, "num_classes": 13},
"kpconv-s3dis-5": {"feature_dimension": 4, "num_classes": 13},
"kpconv-s3dis-6": {"feature_dimension": 4, "num_classes": 13},
}
@staticmethod
def from_pretrained(model_tag, download=True, out_file=None, weight_name="latest", mock_dataset=True):
# Convert inputs to registry format
if PretainedRegistry.MODELS.get(model_tag) is not None:
url = PretainedRegistry.MODELS.get(model_tag)
else:
raise Exception(
"model_tag {} doesn't exist within available models. Here is the list of pre-trained models {}".format(
model_tag, PretainedRegistry.available_models()
)
)
checkpoint_name = model_tag + ".pt"
out_file = os.path.join(CHECKPOINT_DIR, checkpoint_name)
if download:
download_file(url, out_file)
weight_name = weight_name if weight_name is not None else "latest"
checkpoint: ModelCheckpoint = ModelCheckpoint(
CHECKPOINT_DIR, model_tag, weight_name if weight_name is not None else "latest", resume=False,
)
if mock_dataset:
dataset = checkpoint.dataset_properties.copy()
if PretainedRegistry.MOCK_USED_PROPERTIES.get(model_tag) is not None:
for k, v in PretainedRegistry.MOCK_USED_PROPERTIES.get(model_tag).items():
dataset[k] = v
else:
dataset = instantiate_dataset(checkpoint.data_config)
model: BaseModel = checkpoint.create_model(dataset, weight_name=weight_name)
Wandb.set_urls_to_model(model, url)
BaseDataset.set_transform(model, checkpoint.data_config)
return model
@staticmethod
def from_file(path, weight_name="latest", mock_property=None):
"""
Load a pretrained model trained with torch-points3d from file.
return a pretrained model
Parameters
----------
path: str
path of a pretrained model
weight_name: str, optional
name of the weight
mock_property: dict, optional
mock dataset
"""
weight_name = weight_name if weight_name is not None else "latest"
path_dir, name = os.path.split(path)
name = name.split(".")[0] # ModelCheckpoint will add the extension
checkpoint: ModelCheckpoint = ModelCheckpoint(
path_dir, name, weight_name if weight_name is not None else "latest", resume=False,
)
dataset = checkpoint.data_config
if mock_property is not None:
for k, v in mock_property.items():
dataset[k] = v
else:
dataset = instantiate_dataset(checkpoint.data_config)
model: BaseModel = checkpoint.create_model(dataset, weight_name=weight_name)
BaseDataset.set_transform(model, checkpoint.data_config)
return model
@staticmethod
def available_models():
return PretainedRegistry.MODELS.keys()