budui/raycv
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

change

Browse Source
This commit is contained in:
budui
2020-09-05 10:33:35 +08:00
parent 2469bf15fe
commit 39c754374c
21 changed files with 550 additions and 705 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
data/dataset.py
View File

@@ -177,6 +177,13 @@ class GenerationUnpairedDataset(Dataset):
return f"<GenerationUnpairedDataset:\n\tA: {self.A}\n\tB: {self.B}>\nPipeline:\n{self.A.pipeline}"
def normalize_tensor(tensor):
tensor = tensor.float()
tensor -= tensor.min()
tensor /= tensor.max()
return tensor
@DATASET.register_module()
class GenerationUnpairedDatasetWithEdge(Dataset):
def __init__(self, root_a, root_b, random_pair, pipeline, edge_type, edges_path, landmarks_path, size=(256, 256)):
@@ -200,17 +207,19 @@ class GenerationUnpairedDatasetWithEdge(Dataset):
edge_type = self.edge_type
use_landmark = False
edge_path = self.edges_path / f"{op.parent.name}/{op.stem}.{edge_type}.png"
origin_edge = F.to_tensor(Image.open(edge_path).resize(self.size))
origin_edge = F.to_tensor(Image.open(edge_path).resize(self.size, Image.BILINEAR))
if not use_landmark:
return origin_edge
else:
landmark_path = self.landmarks_path / f"{op.parent.name}/{op.stem}.{edge_type}.txt"
landmark_path = self.landmarks_path / f"{op.parent.name}/{op.stem}.txt"
key_points, part_labels, part_edge = dlib_landmark.read_keypoints(landmark_path, size=self.size)
dist_tensor = torch.from_numpy(dlib_landmark.dist_tensor(key_points))
part_labels = torch.from_numpy(part_labels)
edges = origin_edge * (part_labels.sum(0) == 0) # remove edges within face
edges = part_edge + edges
return torch.cat([edges, dist_tensor, part_labels], dim=0)
dist_tensor = normalize_tensor(torch.from_numpy(dlib_landmark.dist_tensor(key_points, size=self.size)))
part_labels = normalize_tensor(torch.from_numpy(part_labels))
part_edge = torch.from_numpy(part_edge).unsqueeze(0).float()
# edges = origin_edge * (part_labels.sum(0) == 0) # remove edges within face
# edges = part_edge + edges
return torch.cat([origin_edge, part_edge, dist_tensor, part_labels])
def __getitem__(self, idx):
a_idx = idx % len(self.A)