DataSet Modifications

This commit is contained in:
Si11ium
2020-07-01 14:15:26 +02:00
parent 6b28519e58
commit 3c1202d5b6
4 changed files with 15 additions and 21 deletions
+5 -6
View File
@@ -111,7 +111,7 @@ class CustomShapeNet(InMemoryDataset):
os.remove(filepath)
try:
config_path.unlink()
except:
except FileNotFoundError:
pass
print('Processed Location "Refreshed" (We deleted the Files)')
except FileNotFoundError:
@@ -138,7 +138,7 @@ class CustomShapeNet(InMemoryDataset):
return data, slices
def _build_config(self):
conf_dict = {key:str(val) for key, val in self.__dict__.items() if '__' not in key and key not in [
conf_dict = {key: str(val) for key, val in self.__dict__.items() if '__' not in key and key not in [
'classes', 'refresh', 'transform', 'data', 'slices'
]}
return conf_dict
@@ -194,10 +194,9 @@ class CustomShapeNet(InMemoryDataset):
continue
tensor = tensor.unsqueeze(0)
if self.poly_as_plane:
tensor[:, -2][tensor[:, -2] == float(self.classes.Plane)] = 4.0
tensor[:, -2][tensor[:, -2] == float(self.classes.Box)] = 4.0
tensor[:, -2][tensor[:, -2] == float(self.classes.Polytope)] = 4.0
tensor[:, -2][tensor[:, -2] == self.classes.Torus] = 3.0
tensor[:, -2][tensor[:, -2] == float(self.classes.Plane)] = 2.0
tensor[:, -2][tensor[:, -2] == float(self.classes.Box)] = 2.0
tensor[:, -2][tensor[:, -2] == float(self.classes.Polytope)] = 2.0
src[key] = tensor
for key, values in src.items():