data_loader2D.py
# Python std.
import os
import random
# 3rd party.
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import helpers
# Project files.
from helpers import load_obj, sample_mesh, add_neighbours, apply_random_rotation
def normalize(pc, N):
minX = torch.min(pc[0, :, :])
minY = torch.min(pc[1, :, :])
minZ = torch.min(pc[2, :, :])
pc = torch.cat(
((pc[0, :, :] - minX).unsqueeze(0), (pc[1, :, :] - minY).unsqueeze(0), (pc[2, :, :] - minZ).unsqueeze(0)), 0)
diff_x = torch.cat((torch.ones(3, N, 1), pc), 2) - \
torch.cat((pc, torch.ones(3, N, 1)), 2)
diff_x_loss = diff_x[:, :, 1:-1].norm(p=2, dim=0).mean()
diff_y = torch.cat((torch.ones(3, 1, N), pc), 1) - \
torch.cat((pc, torch.ones(3, 1, N)), 1)
diff_y_loss = diff_y[:, 1:-1, :].norm(p=2, dim=0).mean()
diff = 0.5 * (diff_y_loss + diff_x_loss)
pc = pc / (N * diff)
return pc
class SF_Dataset(Dataset):
def __init__(self, train=True):
super(SF_Dataset, self).__init__()
self.datapath = []
self.data = {}
self.masks = {}
self.n_masks = 5000
self.N = 64
if train:
data_path = "./ScarfFolds/train"
else:
self.n_masks = 500
data_path = "./ScarfFolds/test"
fns_pc = sorted(os.listdir(data_path))
fns = [val for val in fns_pc]
print('Files ' + str(len(fns)))
if len(fns) == 0:
print("ERRROR : folder is empty")
for fn in fns:
self.datapath.append(os.path.join(data_path, fn))
# Masks
for i in range(self.n_masks):
N = 64
if i % 50 == 0:
print(f"Preprocessing masks: {i + 1}/{self.n_masks}", end='\r')
mask = torch.ones(N, N)
source_x, source_y = random.randint(15, N - 1 - 15), random.randint(15, N - 1 - 15)
type = random.uniform(0.0, 1.0)
if type < 0.5:
# Circle
radius_small = random.randint(20, 40) #random.randint(10, 20)
radius_big = random.randint(50, 70) #random.randint(25, 35)
for a in range(N):
for b in range(N):
if (a - source_x)**2 + (b - source_y)**2 < radius_small**2 or (a - source_x)**2 + (b - source_y)**2 > radius_big**2:
mask[a, b] = 0
else:
# Classic hole
n_missing = random.randint(1200, 2400) #2400, 4800
mask[source_x, source_y] = 0
visited = set()
neighbours = set()
visited.add((source_x, source_y))
add_neighbours(source_x, source_y, N, neighbours, visited)
for j in range(n_missing):
x, y = random.choice(tuple(neighbours))
visited.add((x, y))
neighbours.remove((x, y))
mask[x, y] = 0
add_neighbours(x, y, N, neighbours, visited)
if random.uniform(0.0, 1.0) > 0.5:
mask[source_x, source_y] = 1
self.masks[i] = mask.numpy()
print()
# Cache data
for i, path in enumerate(self.datapath):
if i % 50 == 0:
print(f"Caching data: {i + 1}/{len(self.datapath)}", end='\r')
mesh, tri, normals = load_obj(path)
pc_gt = sample_mesh(mesh, tri, normals, N)
normals_gt = torch.zeros(N*N, 3)
pc_gt = pc_gt.transpose(0, 1).view(3, N, N)
normals_gt = normals_gt.transpose(0, 1).view(3, N, N)
pc_gt = apply_random_rotation(pc_gt)
data = normalize(pc_gt, N)
self.data[i] = (data, normals_gt)
print()
def __getitem__(self, index):
pc_gt, normals_gt = self.data[index]
mask = random.choice(self.masks)
im_input = (pc_gt * mask).detach()
normals_input = (normals_gt * mask).detach()
sample = {'pcloud': pc_gt, 'normals_gt': normals_gt, 'normals_input': normals_input,
'input': im_input, 'mask': mask}
return sample
def __len__(self):
return len(self.datapath)