diff --git a/AdaptiveWingLoss/__pycache__/aux.cpython-310.pyc b/AdaptiveWingLoss/__pycache__/aux.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e28039557364b7f54e55322f5bb07fd52e1c92b3
Binary files /dev/null and b/AdaptiveWingLoss/__pycache__/aux.cpython-310.pyc differ
diff --git a/AdaptiveWingLoss/aux.py b/AdaptiveWingLoss/aux.py
new file mode 100644
index 0000000000000000000000000000000000000000..f566bf532405bdaeb350e7b50dcffb4d328835c3
--- /dev/null
+++ b/AdaptiveWingLoss/aux.py
@@ -0,0 +1,4 @@
+def detect_landmarks(inputs, model_ft):
+ outputs, _ = model_ft(inputs)
+ pred_heatmap = outputs[-1][:, :-1, :, :]
+ return pred_heatmap[:, 96, :, :], pred_heatmap[:, 97, :, :]
diff --git a/AdaptiveWingLoss/core/__init__.py b/AdaptiveWingLoss/core/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/AdaptiveWingLoss/core/__pycache__/__init__.cpython-310.pyc b/AdaptiveWingLoss/core/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5b8d178164b7898d4c7e6b6acbadd1db32e19e41
Binary files /dev/null and b/AdaptiveWingLoss/core/__pycache__/__init__.cpython-310.pyc differ
diff --git a/AdaptiveWingLoss/core/__pycache__/coord_conv.cpython-310.pyc b/AdaptiveWingLoss/core/__pycache__/coord_conv.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d57e2e63717668e6923664247a43c5936336a5cf
Binary files /dev/null and b/AdaptiveWingLoss/core/__pycache__/coord_conv.cpython-310.pyc differ
diff --git a/AdaptiveWingLoss/core/__pycache__/models.cpython-310.pyc b/AdaptiveWingLoss/core/__pycache__/models.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4e8b11af1545643eada6b1248255f34a014f0d60
Binary files /dev/null and b/AdaptiveWingLoss/core/__pycache__/models.cpython-310.pyc differ
diff --git a/AdaptiveWingLoss/core/coord_conv.py b/AdaptiveWingLoss/core/coord_conv.py
new file mode 100755
index 0000000000000000000000000000000000000000..37cae2d71b1c0211534d64a8da4c28530b13efe4
--- /dev/null
+++ b/AdaptiveWingLoss/core/coord_conv.py
@@ -0,0 +1,143 @@
+import torch
+import torch.nn as nn
+
+
+class AddCoordsTh(nn.Module):
+ def __init__(self, x_dim=64, y_dim=64, with_r=False, with_boundary=False):
+ super(AddCoordsTh, self).__init__()
+ self.x_dim = x_dim
+ self.y_dim = y_dim
+ self.with_r = with_r
+ self.with_boundary = with_boundary
+
+ def forward(self, input_tensor, heatmap=None):
+ """
+ input_tensor: (batch, c, x_dim, y_dim)
+ """
+ batch_size_tensor = input_tensor.shape[0]
+
+ xx_ones = torch.ones([1, self.y_dim], dtype=torch.int32).to(input_tensor.device)
+ xx_ones = xx_ones.unsqueeze(-1)
+
+ xx_range = torch.arange(self.x_dim, dtype=torch.int32).unsqueeze(0).to(input_tensor.device)
+ xx_range = xx_range.unsqueeze(1)
+
+ xx_channel = torch.matmul(xx_ones.float(), xx_range.float())
+ xx_channel = xx_channel.unsqueeze(-1)
+
+ yy_ones = torch.ones([1, self.x_dim], dtype=torch.int32).to(input_tensor.device)
+ yy_ones = yy_ones.unsqueeze(1)
+
+ yy_range = torch.arange(self.y_dim, dtype=torch.int32).unsqueeze(0).to(input_tensor.device)
+ yy_range = yy_range.unsqueeze(-1)
+
+ yy_channel = torch.matmul(yy_range.float(), yy_ones.float())
+ yy_channel = yy_channel.unsqueeze(-1)
+
+ xx_channel = xx_channel.permute(0, 3, 2, 1)
+ yy_channel = yy_channel.permute(0, 3, 2, 1)
+
+ xx_channel = xx_channel / (self.x_dim - 1)
+ yy_channel = yy_channel / (self.y_dim - 1)
+
+ xx_channel = xx_channel * 2 - 1
+ yy_channel = yy_channel * 2 - 1
+
+ xx_channel = xx_channel.repeat(batch_size_tensor, 1, 1, 1)
+ yy_channel = yy_channel.repeat(batch_size_tensor, 1, 1, 1)
+
+ if self.with_boundary and type(heatmap) != type(None):
+ boundary_channel = torch.clamp(heatmap[:, -1:, :, :], 0.0, 1.0)
+
+ zero_tensor = torch.zeros_like(xx_channel)
+ xx_boundary_channel = torch.where(boundary_channel > 0.05, xx_channel, zero_tensor)
+ yy_boundary_channel = torch.where(boundary_channel > 0.05, yy_channel, zero_tensor)
+ if self.with_boundary and type(heatmap) != type(None):
+ xx_boundary_channel = xx_boundary_channel.to(input_tensor.device)
+ yy_boundary_channel = yy_boundary_channel.to(input_tensor.device)
+ ret = torch.cat([input_tensor, xx_channel, yy_channel], dim=1)
+
+ if self.with_r:
+ rr = torch.sqrt(torch.pow(xx_channel, 2) + torch.pow(yy_channel, 2))
+ rr = rr / torch.max(rr)
+ ret = torch.cat([ret, rr], dim=1)
+
+ if self.with_boundary and type(heatmap) != type(None):
+ ret = torch.cat([ret, xx_boundary_channel, yy_boundary_channel], dim=1)
+ return ret
+
+
+class CoordConvTh(nn.Module):
+ """CoordConv layer as in the paper."""
+
+ def __init__(self, x_dim, y_dim, with_r, with_boundary, in_channels, first_one=False, *args, **kwargs):
+ super(CoordConvTh, self).__init__()
+ self.addcoords = AddCoordsTh(x_dim=x_dim, y_dim=y_dim, with_r=with_r, with_boundary=with_boundary)
+ in_channels += 2
+ if with_r:
+ in_channels += 1
+ if with_boundary and not first_one:
+ in_channels += 2
+ self.conv = nn.Conv2d(in_channels=in_channels, *args, **kwargs)
+
+ def forward(self, input_tensor, heatmap=None):
+ ret = self.addcoords(input_tensor, heatmap)
+ last_channel = ret[:, -2:, :, :]
+ ret = self.conv(ret)
+ return ret, last_channel
+
+
+"""
+An alternative implementation for PyTorch with auto-infering the x-y dimensions.
+"""
+
+
+class AddCoords(nn.Module):
+ def __init__(self, with_r=False):
+ super().__init__()
+ self.with_r = with_r
+
+ def forward(self, input_tensor):
+ """
+ Args:
+ input_tensor: shape(batch, channel, x_dim, y_dim)
+ """
+ batch_size, _, x_dim, y_dim = input_tensor.size()
+
+ xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1)
+ yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1).transpose(1, 2)
+
+ xx_channel = xx_channel / (x_dim - 1)
+ yy_channel = yy_channel / (y_dim - 1)
+
+ xx_channel = xx_channel * 2 - 1
+ yy_channel = yy_channel * 2 - 1
+
+ xx_channel = xx_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3)
+ yy_channel = yy_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3)
+
+ if input_tensor.is_cuda:
+ xx_channel = xx_channel.to(input_tensor.device)
+ yy_channel = yy_channel.to(input_tensor.device)
+
+ ret = torch.cat([input_tensor, xx_channel.type_as(input_tensor), yy_channel.type_as(input_tensor)], dim=1)
+
+ if self.with_r:
+ rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2))
+ if input_tensor.is_cuda:
+ rr = rr.to(input_tensor.device)
+ ret = torch.cat([ret, rr], dim=1)
+
+ return ret
+
+
+class CoordConv(nn.Module):
+ def __init__(self, in_channels, out_channels, with_r=False, **kwargs):
+ super().__init__()
+ self.addcoords = AddCoords(with_r=with_r)
+ self.conv = nn.Conv2d(in_channels + 2, out_channels, **kwargs)
+
+ def forward(self, x):
+ ret = self.addcoords(x)
+ ret = self.conv(ret)
+ return ret
diff --git a/AdaptiveWingLoss/core/dataloader.py b/AdaptiveWingLoss/core/dataloader.py
new file mode 100755
index 0000000000000000000000000000000000000000..b901540a29e1a72abe2e88c4d052791a9dc4ec36
--- /dev/null
+++ b/AdaptiveWingLoss/core/dataloader.py
@@ -0,0 +1,350 @@
+import copy
+import glob
+import math
+import os
+import random
+import sys
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy.io as sio
+import torch
+from imgaug import augmenters as iaa
+from PIL import Image
+from scipy import interpolate
+from skimage import io
+from skimage import transform as ski_transform
+from skimage.color import rgb2gray
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from torchvision import transforms
+from torchvision import utils
+from torchvision.transforms import Compose
+from torchvision.transforms import Lambda
+from torchvision.transforms.functional import adjust_brightness
+from torchvision.transforms.functional import adjust_contrast
+from torchvision.transforms.functional import adjust_hue
+from torchvision.transforms.functional import adjust_saturation
+
+from utils.utils import cv_crop
+from utils.utils import cv_rotate
+from utils.utils import draw_gaussian
+from utils.utils import fig2data
+from utils.utils import generate_weight_map
+from utils.utils import power_transform
+from utils.utils import shuffle_lr
+from utils.utils import transform
+
+
+class AddBoundary(object):
+ def __init__(self, num_landmarks=68):
+ self.num_landmarks = num_landmarks
+
+ def __call__(self, sample):
+ landmarks_64 = np.floor(sample["landmarks"] / 4.0)
+ if self.num_landmarks == 68:
+ boundaries = {}
+ boundaries["cheek"] = landmarks_64[0:17]
+ boundaries["left_eyebrow"] = landmarks_64[17:22]
+ boundaries["right_eyebrow"] = landmarks_64[22:27]
+ boundaries["uper_left_eyelid"] = landmarks_64[36:40]
+ boundaries["lower_left_eyelid"] = np.array([landmarks_64[i] for i in [36, 41, 40, 39]])
+ boundaries["upper_right_eyelid"] = landmarks_64[42:46]
+ boundaries["lower_right_eyelid"] = np.array([landmarks_64[i] for i in [42, 47, 46, 45]])
+ boundaries["noise"] = landmarks_64[27:31]
+ boundaries["noise_bot"] = landmarks_64[31:36]
+ boundaries["upper_outer_lip"] = landmarks_64[48:55]
+ boundaries["upper_inner_lip"] = np.array([landmarks_64[i] for i in [60, 61, 62, 63, 64]])
+ boundaries["lower_outer_lip"] = np.array([landmarks_64[i] for i in [48, 59, 58, 57, 56, 55, 54]])
+ boundaries["lower_inner_lip"] = np.array([landmarks_64[i] for i in [60, 67, 66, 65, 64]])
+ elif self.num_landmarks == 98:
+ boundaries = {}
+ boundaries["cheek"] = landmarks_64[0:33]
+ boundaries["left_eyebrow"] = landmarks_64[33:38]
+ boundaries["right_eyebrow"] = landmarks_64[42:47]
+ boundaries["uper_left_eyelid"] = landmarks_64[60:65]
+ boundaries["lower_left_eyelid"] = np.array([landmarks_64[i] for i in [60, 67, 66, 65, 64]])
+ boundaries["upper_right_eyelid"] = landmarks_64[68:73]
+ boundaries["lower_right_eyelid"] = np.array([landmarks_64[i] for i in [68, 75, 74, 73, 72]])
+ boundaries["noise"] = landmarks_64[51:55]
+ boundaries["noise_bot"] = landmarks_64[55:60]
+ boundaries["upper_outer_lip"] = landmarks_64[76:83]
+ boundaries["upper_inner_lip"] = np.array([landmarks_64[i] for i in [88, 89, 90, 91, 92]])
+ boundaries["lower_outer_lip"] = np.array([landmarks_64[i] for i in [76, 87, 86, 85, 84, 83, 82]])
+ boundaries["lower_inner_lip"] = np.array([landmarks_64[i] for i in [88, 95, 94, 93, 92]])
+ elif self.num_landmarks == 19:
+ boundaries = {}
+ boundaries["left_eyebrow"] = landmarks_64[0:3]
+ boundaries["right_eyebrow"] = landmarks_64[3:5]
+ boundaries["left_eye"] = landmarks_64[6:9]
+ boundaries["right_eye"] = landmarks_64[9:12]
+ boundaries["noise"] = landmarks_64[12:15]
+
+ elif self.num_landmarks == 29:
+ boundaries = {}
+ boundaries["upper_left_eyebrow"] = np.stack([landmarks_64[0], landmarks_64[4], landmarks_64[2]], axis=0)
+ boundaries["lower_left_eyebrow"] = np.stack([landmarks_64[0], landmarks_64[5], landmarks_64[2]], axis=0)
+ boundaries["upper_right_eyebrow"] = np.stack([landmarks_64[1], landmarks_64[6], landmarks_64[3]], axis=0)
+ boundaries["lower_right_eyebrow"] = np.stack([landmarks_64[1], landmarks_64[7], landmarks_64[3]], axis=0)
+ boundaries["upper_left_eye"] = np.stack([landmarks_64[8], landmarks_64[12], landmarks_64[10]], axis=0)
+ boundaries["lower_left_eye"] = np.stack([landmarks_64[8], landmarks_64[13], landmarks_64[10]], axis=0)
+ boundaries["upper_right_eye"] = np.stack([landmarks_64[9], landmarks_64[14], landmarks_64[11]], axis=0)
+ boundaries["lower_right_eye"] = np.stack([landmarks_64[9], landmarks_64[15], landmarks_64[11]], axis=0)
+ boundaries["noise"] = np.stack([landmarks_64[18], landmarks_64[21], landmarks_64[19]], axis=0)
+ boundaries["outer_upper_lip"] = np.stack([landmarks_64[22], landmarks_64[24], landmarks_64[23]], axis=0)
+ boundaries["inner_upper_lip"] = np.stack([landmarks_64[22], landmarks_64[25], landmarks_64[23]], axis=0)
+ boundaries["outer_lower_lip"] = np.stack([landmarks_64[22], landmarks_64[26], landmarks_64[23]], axis=0)
+ boundaries["inner_lower_lip"] = np.stack([landmarks_64[22], landmarks_64[27], landmarks_64[23]], axis=0)
+ functions = {}
+
+ for key, points in boundaries.items():
+ temp = points[0]
+ new_points = points[0:1, :]
+ for point in points[1:]:
+ if point[0] == temp[0] and point[1] == temp[1]:
+ continue
+ else:
+ new_points = np.concatenate((new_points, np.expand_dims(point, 0)), axis=0)
+ temp = point
+ points = new_points
+ if points.shape[0] == 1:
+ points = np.concatenate((points, points + 0.001), axis=0)
+ k = min(4, points.shape[0])
+ functions[key] = interpolate.splprep([points[:, 0], points[:, 1]], k=k - 1, s=0)
+
+ boundary_map = np.zeros((64, 64))
+
+ fig = plt.figure(figsize=[64 / 96.0, 64 / 96.0], dpi=96)
+
+ ax = fig.add_axes([0, 0, 1, 1])
+
+ ax.axis("off")
+
+ ax.imshow(boundary_map, interpolation="nearest", cmap="gray")
+ # ax.scatter(landmarks[:, 0], landmarks[:, 1], s=1, marker=',', c='w')
+
+ for key in functions.keys():
+ xnew = np.arange(0, 1, 0.01)
+ out = interpolate.splev(xnew, functions[key][0], der=0)
+ plt.plot(out[0], out[1], ",", linewidth=1, color="w")
+
+ img = fig2data(fig)
+
+ plt.close()
+
+ sigma = 1
+ temp = 255 - img[:, :, 1]
+ temp = cv2.distanceTransform(temp, cv2.DIST_L2, cv2.DIST_MASK_PRECISE)
+ temp = temp.astype(np.float32)
+ temp = np.where(temp < 3 * sigma, np.exp(-(temp * temp) / (2 * sigma * sigma)), 0)
+
+ fig = plt.figure(figsize=[64 / 96.0, 64 / 96.0], dpi=96)
+
+ ax = fig.add_axes([0, 0, 1, 1])
+
+ ax.axis("off")
+ ax.imshow(temp, cmap="gray")
+ plt.close()
+
+ boundary_map = fig2data(fig)
+
+ sample["boundary"] = boundary_map[:, :, 0]
+
+ return sample
+
+
+class AddWeightMap(object):
+ def __call__(self, sample):
+ heatmap = sample["heatmap"]
+ boundary = sample["boundary"]
+ heatmap = np.concatenate((heatmap, np.expand_dims(boundary, axis=0)), 0)
+ weight_map = np.zeros_like(heatmap)
+ for i in range(heatmap.shape[0]):
+ weight_map[i] = generate_weight_map(weight_map[i], heatmap[i])
+ sample["weight_map"] = weight_map
+ return sample
+
+
+class ToTensor(object):
+ """Convert ndarrays in sample to Tensors."""
+
+ def __call__(self, sample):
+ image, heatmap, landmarks, boundary, weight_map = (
+ sample["image"],
+ sample["heatmap"],
+ sample["landmarks"],
+ sample["boundary"],
+ sample["weight_map"],
+ )
+
+ # swap color axis because
+ # numpy image: H x W x C
+ # torch image: C X H X W
+ if len(image.shape) == 2:
+ image = np.expand_dims(image, axis=2)
+ image_small = np.expand_dims(image_small, axis=2)
+ image = image.transpose((2, 0, 1))
+ boundary = np.expand_dims(boundary, axis=2)
+ boundary = boundary.transpose((2, 0, 1))
+ return {
+ "image": torch.from_numpy(image).float().div(255.0),
+ "heatmap": torch.from_numpy(heatmap).float(),
+ "landmarks": torch.from_numpy(landmarks).float(),
+ "boundary": torch.from_numpy(boundary).float().div(255.0),
+ "weight_map": torch.from_numpy(weight_map).float(),
+ }
+
+
+class FaceLandmarksDataset(Dataset):
+ """Face Landmarks dataset."""
+
+ def __init__(
+ self,
+ img_dir,
+ landmarks_dir,
+ num_landmarks=68,
+ gray_scale=False,
+ detect_face=False,
+ enhance=False,
+ center_shift=0,
+ transform=None,
+ ):
+ """
+ Args:
+ landmark_dir (string): Path to the mat file with landmarks saved.
+ img_dir (string): Directory with all the images.
+ transform (callable, optional): Optional transform to be applied
+ on a sample.
+ """
+ self.img_dir = img_dir
+ self.landmarks_dir = landmarks_dir
+ self.num_lanmdkars = num_landmarks
+ self.transform = transform
+ self.img_names = glob.glob(self.img_dir + "*.jpg") + glob.glob(self.img_dir + "*.png")
+ self.gray_scale = gray_scale
+ self.detect_face = detect_face
+ self.enhance = enhance
+ self.center_shift = center_shift
+ if self.detect_face:
+ self.face_detector = MTCNN(thresh=[0.5, 0.6, 0.7])
+
+ def __len__(self):
+ return len(self.img_names)
+
+ def __getitem__(self, idx):
+ img_name = self.img_names[idx]
+ pil_image = Image.open(img_name)
+ if pil_image.mode != "RGB":
+ # if input is grayscale image, convert it to 3 channel image
+ if self.enhance:
+ pil_image = power_transform(pil_image, 0.5)
+ temp_image = Image.new("RGB", pil_image.size)
+ temp_image.paste(pil_image)
+ pil_image = temp_image
+ image = np.array(pil_image)
+ if self.gray_scale:
+ image = rgb2gray(image)
+ image = np.expand_dims(image, axis=2)
+ image = np.concatenate((image, image, image), axis=2)
+ image = image * 255.0
+ image = image.astype(np.uint8)
+ if not self.detect_face:
+ center = [450 // 2, 450 // 2 + 0]
+ if self.center_shift != 0:
+ center[0] += int(np.random.uniform(-self.center_shift, self.center_shift))
+ center[1] += int(np.random.uniform(-self.center_shift, self.center_shift))
+ scale = 1.8
+ else:
+ detected_faces = self.face_detector.detect_image(image)
+ if len(detected_faces) > 0:
+ box = detected_faces[0]
+ left, top, right, bottom, _ = box
+ center = [right - (right - left) / 2.0, bottom - (bottom - top) / 2.0]
+ center[1] = center[1] - (bottom - top) * 0.12
+ scale = (right - left + bottom - top) / 195.0
+ else:
+ center = [450 // 2, 450 // 2 + 0]
+ scale = 1.8
+ if self.center_shift != 0:
+ shift = self.center * self.center_shift / 450
+ center[0] += int(np.random.uniform(-shift, shift))
+ center[1] += int(np.random.uniform(-shift, shift))
+ base_name = os.path.basename(img_name)
+ landmarks_base_name = base_name[:-4] + "_pts.mat"
+ landmarks_name = os.path.join(self.landmarks_dir, landmarks_base_name)
+ if os.path.isfile(landmarks_name):
+ mat_data = sio.loadmat(landmarks_name)
+ landmarks = mat_data["pts_2d"]
+ elif os.path.isfile(landmarks_name[:-8] + ".pts.npy"):
+ landmarks = np.load(landmarks_name[:-8] + ".pts.npy")
+ else:
+ landmarks = []
+ heatmap = []
+
+ if landmarks != []:
+ new_image, new_landmarks = cv_crop(image, landmarks, center, scale, 256, self.center_shift)
+ tries = 0
+ while self.center_shift != 0 and tries < 5 and (np.max(new_landmarks) > 240 or np.min(new_landmarks) < 15):
+ center = [450 // 2, 450 // 2 + 0]
+ scale += 0.05
+ center[0] += int(np.random.uniform(-self.center_shift, self.center_shift))
+ center[1] += int(np.random.uniform(-self.center_shift, self.center_shift))
+
+ new_image, new_landmarks = cv_crop(image, landmarks, center, scale, 256, self.center_shift)
+ tries += 1
+ if np.max(new_landmarks) > 250 or np.min(new_landmarks) < 5:
+ center = [450 // 2, 450 // 2 + 0]
+ scale = 2.25
+ new_image, new_landmarks = cv_crop(image, landmarks, center, scale, 256, 100)
+ assert np.min(new_landmarks) > 0 and np.max(new_landmarks) < 256, "Landmarks out of boundary!"
+ image = new_image
+ landmarks = new_landmarks
+ heatmap = np.zeros((self.num_lanmdkars, 64, 64))
+ for i in range(self.num_lanmdkars):
+ if landmarks[i][0] > 0:
+ heatmap[i] = draw_gaussian(heatmap[i], landmarks[i] / 4.0 + 1, 1)
+ sample = {"image": image, "heatmap": heatmap, "landmarks": landmarks}
+ if self.transform:
+ sample = self.transform(sample)
+
+ return sample
+
+
+def get_dataset(
+ val_img_dir,
+ val_landmarks_dir,
+ batch_size,
+ num_landmarks=68,
+ rotation=0,
+ scale=0,
+ center_shift=0,
+ random_flip=False,
+ brightness=0,
+ contrast=0,
+ saturation=0,
+ blur=False,
+ noise=False,
+ jpeg_effect=False,
+ random_occlusion=False,
+ gray_scale=False,
+ detect_face=False,
+ enhance=False,
+):
+ val_transforms = transforms.Compose([AddBoundary(num_landmarks), AddWeightMap(), ToTensor()])
+
+ val_dataset = FaceLandmarksDataset(
+ val_img_dir,
+ val_landmarks_dir,
+ num_landmarks=num_landmarks,
+ gray_scale=gray_scale,
+ detect_face=detect_face,
+ enhance=enhance,
+ transform=val_transforms,
+ )
+
+ val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=6)
+ data_loaders = {"val": val_dataloader}
+ dataset_sizes = {}
+ dataset_sizes["val"] = len(val_dataset)
+ return data_loaders, dataset_sizes
diff --git a/AdaptiveWingLoss/core/evaler.py b/AdaptiveWingLoss/core/evaler.py
new file mode 100755
index 0000000000000000000000000000000000000000..9a1a3c26560dc6a34067df513f7cc85798fa3b25
--- /dev/null
+++ b/AdaptiveWingLoss/core/evaler.py
@@ -0,0 +1,125 @@
+import matplotlib
+
+matplotlib.use("Agg")
+import math
+import torch
+import copy
+import time
+from torch.autograd import Variable
+import shutil
+from skimage import io
+import numpy as np
+from utils.utils import fan_NME, show_landmarks, get_preds_fromhm
+from PIL import Image, ImageDraw
+import os
+import sys
+import cv2
+import matplotlib.pyplot as plt
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+def eval_model(
+ model, dataloaders, dataset_sizes, writer, use_gpu=True, epoches=5, dataset="val", save_path="./", num_landmarks=68
+):
+ global_nme = 0
+ model.eval()
+ for epoch in range(epoches):
+ running_loss = 0
+ step = 0
+ total_nme = 0
+ total_count = 0
+ fail_count = 0
+ nmes = []
+ # running_corrects = 0
+
+ # Iterate over data.
+ with torch.no_grad():
+ for data in dataloaders[dataset]:
+ total_runtime = 0
+ run_count = 0
+ step_start = time.time()
+ step += 1
+ # get the inputs
+ inputs = data["image"].type(torch.FloatTensor)
+ labels_heatmap = data["heatmap"].type(torch.FloatTensor)
+ labels_boundary = data["boundary"].type(torch.FloatTensor)
+ landmarks = data["landmarks"].type(torch.FloatTensor)
+ loss_weight_map = data["weight_map"].type(torch.FloatTensor)
+ # wrap them in Variable
+ if use_gpu:
+ inputs = inputs.to(device)
+ labels_heatmap = labels_heatmap.to(device)
+ labels_boundary = labels_boundary.to(device)
+ loss_weight_map = loss_weight_map.to(device)
+ else:
+ inputs, labels_heatmap = Variable(inputs), Variable(labels_heatmap)
+ labels_boundary = Variable(labels_boundary)
+ labels = torch.cat((labels_heatmap, labels_boundary), 1)
+ single_start = time.time()
+ outputs, boundary_channels = model(inputs)
+ single_end = time.time()
+ total_runtime += time.time() - single_start
+ run_count += 1
+ step_end = time.time()
+ for i in range(inputs.shape[0]):
+ img = inputs[i]
+ img = img.cpu().numpy()
+ img = img.transpose((1, 2, 0)) * 255.0
+ img = img.astype(np.uint8)
+ img = Image.fromarray(img)
+ # pred_heatmap = outputs[-1][i].detach().cpu()[:-1, :, :]
+ pred_heatmap = outputs[-1][:, :-1, :, :][i].detach().cpu()
+ pred_landmarks, _ = get_preds_fromhm(pred_heatmap.unsqueeze(0))
+ pred_landmarks = pred_landmarks.squeeze().numpy()
+
+ gt_landmarks = data["landmarks"][i].numpy()
+ if num_landmarks == 68:
+ left_eye = np.average(gt_landmarks[36:42], axis=0)
+ right_eye = np.average(gt_landmarks[42:48], axis=0)
+ norm_factor = np.linalg.norm(left_eye - right_eye)
+ # norm_factor = np.linalg.norm(gt_landmarks[36]- gt_landmarks[45])
+
+ elif num_landmarks == 98:
+ norm_factor = np.linalg.norm(gt_landmarks[60] - gt_landmarks[72])
+ elif num_landmarks == 19:
+ left, top = gt_landmarks[-2, :]
+ right, bottom = gt_landmarks[-1, :]
+ norm_factor = math.sqrt(abs(right - left) * abs(top - bottom))
+ gt_landmarks = gt_landmarks[:-2, :]
+ elif num_landmarks == 29:
+ # norm_factor = np.linalg.norm(gt_landmarks[8]- gt_landmarks[9])
+ norm_factor = np.linalg.norm(gt_landmarks[16] - gt_landmarks[17])
+ single_nme = (
+ np.sum(np.linalg.norm(pred_landmarks * 4 - gt_landmarks, axis=1)) / pred_landmarks.shape[0]
+ ) / norm_factor
+
+ nmes.append(single_nme)
+ total_count += 1
+ if single_nme > 0.1:
+ fail_count += 1
+ if step % 10 == 0:
+ print(
+ "Step {} Time: {:.6f} Input Mean: {:.6f} Output Mean: {:.6f}".format(
+ step, step_end - step_start, torch.mean(labels), torch.mean(outputs[0])
+ )
+ )
+ # gt_landmarks = landmarks.numpy()
+ # pred_heatmap = outputs[-1].to('cpu').numpy()
+ gt_landmarks = landmarks
+ batch_nme = fan_NME(outputs[-1][:, :-1, :, :].detach().cpu(), gt_landmarks, num_landmarks)
+ # batch_nme = 0
+ total_nme += batch_nme
+ epoch_nme = total_nme / dataset_sizes["val"]
+ global_nme += epoch_nme
+ nme_save_path = os.path.join(save_path, "nme_log.npy")
+ np.save(nme_save_path, np.array(nmes))
+ print(
+ "NME: {:.6f} Failure Rate: {:.6f} Total Count: {:.6f} Fail Count: {:.6f}".format(
+ epoch_nme, fail_count / total_count, total_count, fail_count
+ )
+ )
+ print("Evaluation done! Average NME: {:.6f}".format(global_nme / epoches))
+ print("Everage runtime for a single batch: {:.6f}".format(total_runtime / run_count))
+ return model
diff --git a/AdaptiveWingLoss/core/models.py b/AdaptiveWingLoss/core/models.py
new file mode 100755
index 0000000000000000000000000000000000000000..f2c3d5c0616db1869a8ff9c9e74dffaf4c1c8828
--- /dev/null
+++ b/AdaptiveWingLoss/core/models.py
@@ -0,0 +1,239 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from AdaptiveWingLoss.core.coord_conv import CoordConvTh
+
+
+def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False, dilation=1):
+ "3x3 convolution with padding"
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=strd, padding=padding, bias=bias, dilation=dilation)
+
+
+class BasicBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None):
+ super(BasicBlock, self).__init__()
+ self.conv1 = conv3x3(inplanes, planes, stride)
+ # self.bn1 = nn.BatchNorm2d(planes)
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = conv3x3(planes, planes)
+ # self.bn2 = nn.BatchNorm2d(planes)
+ self.downsample = downsample
+ self.stride = stride
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ # out = self.bn1(out)
+ out = self.relu(out)
+
+ out = self.conv2(out)
+ # out = self.bn2(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.relu(out)
+
+ return out
+
+
+class ConvBlock(nn.Module):
+ def __init__(self, in_planes, out_planes):
+ super(ConvBlock, self).__init__()
+ self.bn1 = nn.BatchNorm2d(in_planes)
+ self.conv1 = conv3x3(in_planes, int(out_planes / 2))
+ self.bn2 = nn.BatchNorm2d(int(out_planes / 2))
+ self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4), padding=1, dilation=1)
+ self.bn3 = nn.BatchNorm2d(int(out_planes / 4))
+ self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4), padding=1, dilation=1)
+
+ if in_planes != out_planes:
+ self.downsample = nn.Sequential(
+ nn.BatchNorm2d(in_planes),
+ nn.ReLU(True),
+ nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, bias=False),
+ )
+ else:
+ self.downsample = None
+
+ def forward(self, x):
+ residual = x
+
+ out1 = self.bn1(x)
+ out1 = F.relu(out1, True)
+ out1 = self.conv1(out1)
+
+ out2 = self.bn2(out1)
+ out2 = F.relu(out2, True)
+ out2 = self.conv2(out2)
+
+ out3 = self.bn3(out2)
+ out3 = F.relu(out3, True)
+ out3 = self.conv3(out3)
+
+ out3 = torch.cat((out1, out2, out3), 1)
+
+ if self.downsample is not None:
+ residual = self.downsample(residual)
+
+ out3 += residual
+
+ return out3
+
+
+class HourGlass(nn.Module):
+ def __init__(self, num_modules, depth, num_features, first_one=False):
+ super(HourGlass, self).__init__()
+ self.num_modules = num_modules
+ self.depth = depth
+ self.features = num_features
+ self.coordconv = CoordConvTh(
+ x_dim=64,
+ y_dim=64,
+ with_r=True,
+ with_boundary=True,
+ in_channels=256,
+ first_one=first_one,
+ out_channels=256,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+ self._generate_network(self.depth)
+
+ def _generate_network(self, level):
+ self.add_module("b1_" + str(level), ConvBlock(256, 256))
+
+ self.add_module("b2_" + str(level), ConvBlock(256, 256))
+
+ if level > 1:
+ self._generate_network(level - 1)
+ else:
+ self.add_module("b2_plus_" + str(level), ConvBlock(256, 256))
+
+ self.add_module("b3_" + str(level), ConvBlock(256, 256))
+
+ def _forward(self, level, inp):
+ # Upper branch
+ up1 = inp
+ up1 = self._modules["b1_" + str(level)](up1)
+
+ # Lower branch
+ low1 = F.avg_pool2d(inp, 2, stride=2)
+ low1 = self._modules["b2_" + str(level)](low1)
+
+ if level > 1:
+ low2 = self._forward(level - 1, low1)
+ else:
+ low2 = low1
+ low2 = self._modules["b2_plus_" + str(level)](low2)
+
+ low3 = low2
+ low3 = self._modules["b3_" + str(level)](low3)
+
+ up2 = F.upsample(low3, scale_factor=2, mode="nearest")
+
+ return up1 + up2
+
+ def forward(self, x, heatmap):
+ x, last_channel = self.coordconv(x, heatmap)
+ return self._forward(self.depth, x), last_channel
+
+
+class FAN(nn.Module):
+ def __init__(self, num_modules=1, end_relu=False, gray_scale=False, num_landmarks=68):
+ super(FAN, self).__init__()
+ self.num_modules = num_modules
+ self.gray_scale = gray_scale
+ self.end_relu = end_relu
+ self.num_landmarks = num_landmarks
+
+ # Base part
+ if self.gray_scale:
+ self.conv1 = CoordConvTh(
+ x_dim=256,
+ y_dim=256,
+ with_r=True,
+ with_boundary=False,
+ in_channels=3,
+ out_channels=64,
+ kernel_size=7,
+ stride=2,
+ padding=3,
+ )
+ else:
+ self.conv1 = CoordConvTh(
+ x_dim=256,
+ y_dim=256,
+ with_r=True,
+ with_boundary=False,
+ in_channels=3,
+ out_channels=64,
+ kernel_size=7,
+ stride=2,
+ padding=3,
+ )
+ self.bn1 = nn.BatchNorm2d(64)
+ self.conv2 = ConvBlock(64, 128)
+ self.conv3 = ConvBlock(128, 128)
+ self.conv4 = ConvBlock(128, 256)
+
+ # Stacking part
+ for hg_module in range(self.num_modules):
+ if hg_module == 0:
+ first_one = True
+ else:
+ first_one = False
+ self.add_module("m" + str(hg_module), HourGlass(1, 4, 256, first_one))
+ self.add_module("top_m_" + str(hg_module), ConvBlock(256, 256))
+ self.add_module("conv_last" + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+ self.add_module("bn_end" + str(hg_module), nn.BatchNorm2d(256))
+ self.add_module("l" + str(hg_module), nn.Conv2d(256, num_landmarks + 1, kernel_size=1, stride=1, padding=0))
+
+ if hg_module < self.num_modules - 1:
+ self.add_module("bl" + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+ self.add_module(
+ "al" + str(hg_module), nn.Conv2d(num_landmarks + 1, 256, kernel_size=1, stride=1, padding=0)
+ )
+
+ def forward(self, x):
+ x, _ = self.conv1(x)
+ x = F.relu(self.bn1(x), True)
+ # x = F.relu(self.bn1(self.conv1(x)), True)
+ x = F.avg_pool2d(self.conv2(x), 2, stride=2)
+ x = self.conv3(x)
+ x = self.conv4(x)
+
+ previous = x
+
+ outputs = []
+ boundary_channels = []
+ tmp_out = None
+ for i in range(self.num_modules):
+ hg, boundary_channel = self._modules["m" + str(i)](previous, tmp_out)
+
+ ll = hg
+ ll = self._modules["top_m_" + str(i)](ll)
+
+ ll = F.relu(self._modules["bn_end" + str(i)](self._modules["conv_last" + str(i)](ll)), True)
+
+ # Predict heatmaps
+ tmp_out = self._modules["l" + str(i)](ll)
+ if self.end_relu:
+ tmp_out = F.relu(tmp_out) # HACK: Added relu
+ outputs.append(tmp_out)
+ boundary_channels.append(boundary_channel)
+
+ if i < self.num_modules - 1:
+ ll = self._modules["bl" + str(i)](ll)
+ tmp_out_ = self._modules["al" + str(i)](tmp_out)
+ previous = previous + ll + tmp_out_
+
+ return outputs, boundary_channels
diff --git a/AdaptiveWingLoss/utils/__init__.py b/AdaptiveWingLoss/utils/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/AdaptiveWingLoss/utils/utils.py b/AdaptiveWingLoss/utils/utils.py
new file mode 100755
index 0000000000000000000000000000000000000000..9daa0be2b71c5e7d92679f38692a891ac7e3cc0a
--- /dev/null
+++ b/AdaptiveWingLoss/utils/utils.py
@@ -0,0 +1,437 @@
+from __future__ import division
+from __future__ import print_function
+
+import math
+import os
+import sys
+
+import cv2
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from PIL import Image
+from scipy import ndimage
+from skimage import io
+from skimage import transform as ski_transform
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from torchvision import transforms
+from torchvision import utils
+
+
+def _gaussian(
+ size=3,
+ sigma=0.25,
+ amplitude=1,
+ normalize=False,
+ width=None,
+ height=None,
+ sigma_horz=None,
+ sigma_vert=None,
+ mean_horz=0.5,
+ mean_vert=0.5,
+):
+ # handle some defaults
+ if width is None:
+ width = size
+ if height is None:
+ height = size
+ if sigma_horz is None:
+ sigma_horz = sigma
+ if sigma_vert is None:
+ sigma_vert = sigma
+ center_x = mean_horz * width + 0.5
+ center_y = mean_vert * height + 0.5
+ gauss = np.empty((height, width), dtype=np.float32)
+ # generate kernel
+ for i in range(height):
+ for j in range(width):
+ gauss[i][j] = amplitude * math.exp(
+ -(
+ math.pow((j + 1 - center_x) / (sigma_horz * width), 2) / 2.0
+ + math.pow((i + 1 - center_y) / (sigma_vert * height), 2) / 2.0
+ )
+ )
+ if normalize:
+ gauss = gauss / np.sum(gauss)
+ return gauss
+
+
+def draw_gaussian(image, point, sigma):
+ # Check if the gaussian is inside
+ ul = [np.floor(np.floor(point[0]) - 3 * sigma), np.floor(np.floor(point[1]) - 3 * sigma)]
+ br = [np.floor(np.floor(point[0]) + 3 * sigma), np.floor(np.floor(point[1]) + 3 * sigma)]
+ if ul[0] > image.shape[1] or ul[1] > image.shape[0] or br[0] < 1 or br[1] < 1:
+ return image
+ size = 6 * sigma + 1
+ g = _gaussian(size)
+ g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) - int(max(1, ul[0])) + int(max(1, -ul[0]))]
+ g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) - int(max(1, ul[1])) + int(max(1, -ul[1]))]
+ img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]
+ img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]
+ assert g_x[0] > 0 and g_y[1] > 0
+ correct = False
+ while not correct:
+ try:
+ image[img_y[0] - 1 : img_y[1], img_x[0] - 1 : img_x[1]] = (
+ image[img_y[0] - 1 : img_y[1], img_x[0] - 1 : img_x[1]] + g[g_y[0] - 1 : g_y[1], g_x[0] - 1 : g_x[1]]
+ )
+ correct = True
+ except:
+ print(
+ "img_x: {}, img_y: {}, g_x:{}, g_y:{}, point:{}, g_shape:{}, ul:{}, br:{}".format(
+ img_x, img_y, g_x, g_y, point, g.shape, ul, br
+ )
+ )
+ ul = [np.floor(np.floor(point[0]) - 3 * sigma), np.floor(np.floor(point[1]) - 3 * sigma)]
+ br = [np.floor(np.floor(point[0]) + 3 * sigma), np.floor(np.floor(point[1]) + 3 * sigma)]
+ g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) - int(max(1, ul[0])) + int(max(1, -ul[0]))]
+ g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) - int(max(1, ul[1])) + int(max(1, -ul[1]))]
+ img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))]
+ img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))]
+ pass
+ image[image > 1] = 1
+ return image
+
+
+def transform(point, center, scale, resolution, rotation=0, invert=False):
+ _pt = np.ones(3)
+ _pt[0] = point[0]
+ _pt[1] = point[1]
+
+ h = 200.0 * scale
+ t = np.eye(3)
+ t[0, 0] = resolution / h
+ t[1, 1] = resolution / h
+ t[0, 2] = resolution * (-center[0] / h + 0.5)
+ t[1, 2] = resolution * (-center[1] / h + 0.5)
+
+ if rotation != 0:
+ rotation = -rotation
+ r = np.eye(3)
+ ang = rotation * math.pi / 180.0
+ s = math.sin(ang)
+ c = math.cos(ang)
+ r[0][0] = c
+ r[0][1] = -s
+ r[1][0] = s
+ r[1][1] = c
+
+ t_ = np.eye(3)
+ t_[0][2] = -resolution / 2.0
+ t_[1][2] = -resolution / 2.0
+ t_inv = torch.eye(3)
+ t_inv[0][2] = resolution / 2.0
+ t_inv[1][2] = resolution / 2.0
+ t = reduce(np.matmul, [t_inv, r, t_, t])
+
+ if invert:
+ t = np.linalg.inv(t)
+ new_point = (np.matmul(t, _pt))[0:2]
+
+ return new_point.astype(int)
+
+
+def cv_crop(image, landmarks, center, scale, resolution=256, center_shift=0):
+ new_image = cv2.copyMakeBorder(
+ image, center_shift, center_shift, center_shift, center_shift, cv2.BORDER_CONSTANT, value=[0, 0, 0]
+ )
+ new_landmarks = landmarks.copy()
+ if center_shift != 0:
+ center[0] += center_shift
+ center[1] += center_shift
+ new_landmarks = new_landmarks + center_shift
+ length = 200 * scale
+ top = int(center[1] - length // 2)
+ bottom = int(center[1] + length // 2)
+ left = int(center[0] - length // 2)
+ right = int(center[0] + length // 2)
+ y_pad = abs(min(top, new_image.shape[0] - bottom, 0))
+ x_pad = abs(min(left, new_image.shape[1] - right, 0))
+ top, bottom, left, right = top + y_pad, bottom + y_pad, left + x_pad, right + x_pad
+ new_image = cv2.copyMakeBorder(new_image, y_pad, y_pad, x_pad, x_pad, cv2.BORDER_CONSTANT, value=[0, 0, 0])
+ new_image = new_image[top:bottom, left:right]
+ new_image = cv2.resize(new_image, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR)
+ new_landmarks[:, 0] = (new_landmarks[:, 0] + x_pad - left) * resolution / length
+ new_landmarks[:, 1] = (new_landmarks[:, 1] + y_pad - top) * resolution / length
+ return new_image, new_landmarks
+
+
+def cv_rotate(image, landmarks, heatmap, rot, scale, resolution=256):
+ img_mat = cv2.getRotationMatrix2D((resolution // 2, resolution // 2), rot, scale)
+ ones = np.ones(shape=(landmarks.shape[0], 1))
+ stacked_landmarks = np.hstack([landmarks, ones])
+ new_landmarks = img_mat.dot(stacked_landmarks.T).T
+ if np.max(new_landmarks) > 255 or np.min(new_landmarks) < 0:
+ return image, landmarks, heatmap
+ else:
+ new_image = cv2.warpAffine(image, img_mat, (resolution, resolution))
+ if heatmap is not None:
+ new_heatmap = np.zeros((heatmap.shape[0], 64, 64))
+ for i in range(heatmap.shape[0]):
+ if new_landmarks[i][0] > 0:
+ new_heatmap[i] = draw_gaussian(new_heatmap[i], new_landmarks[i] / 4.0 + 1, 1)
+ return new_image, new_landmarks, new_heatmap
+
+
+def show_landmarks(image, heatmap, gt_landmarks, gt_heatmap):
+ """Show image with pred_landmarks"""
+ pred_landmarks = []
+ pred_landmarks, _ = get_preds_fromhm(torch.from_numpy(heatmap).unsqueeze(0))
+ pred_landmarks = pred_landmarks.squeeze() * 4
+
+ # pred_landmarks2 = get_preds_fromhm2(heatmap)
+ heatmap = np.max(gt_heatmap, axis=0)
+ heatmap = heatmap / np.max(heatmap)
+ # image = ski_transform.resize(image, (64, 64))*255
+ image = image.astype(np.uint8)
+ heatmap = np.max(gt_heatmap, axis=0)
+ heatmap = ski_transform.resize(heatmap, (image.shape[0], image.shape[1]))
+ heatmap *= 255
+ heatmap = heatmap.astype(np.uint8)
+ heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+ plt.imshow(image)
+ plt.scatter(gt_landmarks[:, 0], gt_landmarks[:, 1], s=0.5, marker=".", c="g")
+ plt.scatter(pred_landmarks[:, 0], pred_landmarks[:, 1], s=0.5, marker=".", c="r")
+ plt.pause(0.001) # pause a bit so that plots are updated
+
+
+def fan_NME(pred_heatmaps, gt_landmarks, num_landmarks=68):
+ """
+ Calculate total NME for a batch of data
+
+ Args:
+ pred_heatmaps: torch tensor of size [batch, points, height, width]
+ gt_landmarks: torch tesnsor of size [batch, points, x, y]
+
+ Returns:
+ nme: sum of nme for this batch
+ """
+ nme = 0
+ pred_landmarks, _ = get_preds_fromhm(pred_heatmaps)
+ pred_landmarks = pred_landmarks.numpy()
+ gt_landmarks = gt_landmarks.numpy()
+ for i in range(pred_landmarks.shape[0]):
+ pred_landmark = pred_landmarks[i] * 4.0
+ gt_landmark = gt_landmarks[i]
+
+ if num_landmarks == 68:
+ left_eye = np.average(gt_landmark[36:42], axis=0)
+ right_eye = np.average(gt_landmark[42:48], axis=0)
+ norm_factor = np.linalg.norm(left_eye - right_eye)
+ # norm_factor = np.linalg.norm(gt_landmark[36]- gt_landmark[45])
+ elif num_landmarks == 98:
+ norm_factor = np.linalg.norm(gt_landmark[60] - gt_landmark[72])
+ elif num_landmarks == 19:
+ left, top = gt_landmark[-2, :]
+ right, bottom = gt_landmark[-1, :]
+ norm_factor = math.sqrt(abs(right - left) * abs(top - bottom))
+ gt_landmark = gt_landmark[:-2, :]
+ elif num_landmarks == 29:
+ # norm_factor = np.linalg.norm(gt_landmark[8]- gt_landmark[9])
+ norm_factor = np.linalg.norm(gt_landmark[16] - gt_landmark[17])
+ nme += (np.sum(np.linalg.norm(pred_landmark - gt_landmark, axis=1)) / pred_landmark.shape[0]) / norm_factor
+ return nme
+
+
+def fan_NME_hm(pred_heatmaps, gt_heatmaps, num_landmarks=68):
+ """
+ Calculate total NME for a batch of data
+
+ Args:
+ pred_heatmaps: torch tensor of size [batch, points, height, width]
+ gt_landmarks: torch tesnsor of size [batch, points, x, y]
+
+ Returns:
+ nme: sum of nme for this batch
+ """
+ nme = 0
+ pred_landmarks, _ = get_index_fromhm(pred_heatmaps)
+ pred_landmarks = pred_landmarks.numpy()
+ gt_landmarks = gt_landmarks.numpy()
+ for i in range(pred_landmarks.shape[0]):
+ pred_landmark = pred_landmarks[i] * 4.0
+ gt_landmark = gt_landmarks[i]
+ if num_landmarks == 68:
+ left_eye = np.average(gt_landmark[36:42], axis=0)
+ right_eye = np.average(gt_landmark[42:48], axis=0)
+ norm_factor = np.linalg.norm(left_eye - right_eye)
+ else:
+ norm_factor = np.linalg.norm(gt_landmark[60] - gt_landmark[72])
+ nme += (np.sum(np.linalg.norm(pred_landmark - gt_landmark, axis=1)) / pred_landmark.shape[0]) / norm_factor
+ return nme
+
+
+def power_transform(img, power):
+ img = np.array(img)
+ img_new = np.power((img / 255.0), power) * 255.0
+ img_new = img_new.astype(np.uint8)
+ img_new = Image.fromarray(img_new)
+ return img_new
+
+
+def get_preds_fromhm(hm, center=None, scale=None, rot=None):
+ max, idx = torch.max(hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
+ idx += 1
+ preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
+ preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)
+ preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)
+
+ for i in range(preds.size(0)):
+ for j in range(preds.size(1)):
+ hm_ = hm[i, j, :]
+ pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
+ if pX > 0 and pX < 63 and pY > 0 and pY < 63:
+ diff = torch.FloatTensor([hm_[pY, pX + 1] - hm_[pY, pX - 1], hm_[pY + 1, pX] - hm_[pY - 1, pX]])
+ preds[i, j].add_(diff.sign_().mul_(0.25))
+
+ preds.add_(-0.5)
+
+ preds_orig = torch.zeros(preds.size())
+ if center is not None and scale is not None:
+ for i in range(hm.size(0)):
+ for j in range(hm.size(1)):
+ preds_orig[i, j] = transform(preds[i, j], center, scale, hm.size(2), rot, True)
+
+ return preds, preds_orig
+
+
+def get_index_fromhm(hm):
+ max, idx = torch.max(hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
+ preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
+ preds[..., 0].remainder_(hm.size(3))
+ preds[..., 1].div_(hm.size(2)).floor_()
+
+ for i in range(preds.size(0)):
+ for j in range(preds.size(1)):
+ hm_ = hm[i, j, :]
+ pX, pY = int(preds[i, j, 0]), int(preds[i, j, 1])
+ if pX > 0 and pX < 63 and pY > 0 and pY < 63:
+ diff = torch.FloatTensor([hm_[pY, pX + 1] - hm_[pY, pX - 1], hm_[pY + 1, pX] - hm_[pY - 1, pX]])
+ preds[i, j].add_(diff.sign_().mul_(0.25))
+
+ return preds
+
+
+def shuffle_lr(parts, num_landmarks=68, pairs=None):
+ if num_landmarks == 68:
+ if pairs is None:
+ pairs = [
+ [0, 16],
+ [1, 15],
+ [2, 14],
+ [3, 13],
+ [4, 12],
+ [5, 11],
+ [6, 10],
+ [7, 9],
+ [17, 26],
+ [18, 25],
+ [19, 24],
+ [20, 23],
+ [21, 22],
+ [36, 45],
+ [37, 44],
+ [38, 43],
+ [39, 42],
+ [41, 46],
+ [40, 47],
+ [31, 35],
+ [32, 34],
+ [50, 52],
+ [49, 53],
+ [48, 54],
+ [61, 63],
+ [60, 64],
+ [67, 65],
+ [59, 55],
+ [58, 56],
+ ]
+ elif num_landmarks == 98:
+ if pairs is None:
+ pairs = [
+ [0, 32],
+ [1, 31],
+ [2, 30],
+ [3, 29],
+ [4, 28],
+ [5, 27],
+ [6, 26],
+ [7, 25],
+ [8, 24],
+ [9, 23],
+ [10, 22],
+ [11, 21],
+ [12, 20],
+ [13, 19],
+ [14, 18],
+ [15, 17],
+ [33, 46],
+ [34, 45],
+ [35, 44],
+ [36, 43],
+ [37, 42],
+ [38, 50],
+ [39, 49],
+ [40, 48],
+ [41, 47],
+ [60, 72],
+ [61, 71],
+ [62, 70],
+ [63, 69],
+ [64, 68],
+ [65, 75],
+ [66, 74],
+ [67, 73],
+ [96, 97],
+ [55, 59],
+ [56, 58],
+ [76, 82],
+ [77, 81],
+ [78, 80],
+ [88, 92],
+ [89, 91],
+ [95, 93],
+ [87, 83],
+ [86, 84],
+ ]
+ elif num_landmarks == 19:
+ if pairs is None:
+ pairs = [[0, 5], [1, 4], [2, 3], [6, 11], [7, 10], [8, 9], [12, 14], [15, 17]]
+ elif num_landmarks == 29:
+ if pairs is None:
+ pairs = [[0, 1], [4, 6], [5, 7], [2, 3], [8, 9], [12, 14], [16, 17], [13, 15], [10, 11], [18, 19], [22, 23]]
+ for matched_p in pairs:
+ idx1, idx2 = matched_p[0], matched_p[1]
+ tmp = np.copy(parts[idx1])
+ np.copyto(parts[idx1], parts[idx2])
+ np.copyto(parts[idx2], tmp)
+ return parts
+
+
+def generate_weight_map(weight_map, heatmap):
+
+ k_size = 3
+ dilate = ndimage.grey_dilation(heatmap, size=(k_size, k_size))
+ weight_map[np.where(dilate > 0.2)] = 1
+ return weight_map
+
+
+def fig2data(fig):
+ """
+ @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it
+ @param fig a matplotlib figure
+ @return a numpy 3D array of RGBA values
+ """
+ # draw the renderer
+ fig.canvas.draw()
+
+ # Get the RGB buffer from the figure
+ w, h = fig.canvas.get_width_height()
+ buf = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
+ buf.shape = (w, h, 3)
+
+ # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode
+ buf = np.roll(buf, 3, axis=2)
+ return buf
diff --git a/Deep3DFaceRecon_pytorch/LICENSE b/Deep3DFaceRecon_pytorch/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..7e2b184589f4acbeb98eecd453c28b41ff551471
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Sicheng Xu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/Deep3DFaceRecon_pytorch/README.md b/Deep3DFaceRecon_pytorch/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..8a1b306c8cdec09533ed4cd27ca0799b7f9ab53f
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/README.md
@@ -0,0 +1,256 @@
+## Accurate 3D Face Reconstruction with Weakly-Supervised Learning: From Single Image to Image Set —— PyTorch implementation ##
+
+
+
+
+
+This is an unofficial official pytorch implementation of the following paper:
+
+Y. Deng, J. Yang, S. Xu, D. Chen, Y. Jia, and X. Tong, [Accurate 3D Face Reconstruction with Weakly-Supervised Learning: From Single Image to Image Set](https://arxiv.org/abs/1903.08527), IEEE Computer Vision and Pattern Recognition Workshop (CVPRW) on Analysis and Modeling of Faces and Gestures (AMFG), 2019. (**_Best Paper Award!_**)
+
+The method enforces a hybrid-level weakly-supervised training for CNN-based 3D face reconstruction. It is fast, accurate, and robust to pose and occlussions. It achieves state-of-the-art performance on multiple datasets such as FaceWarehouse, MICC Florence and NoW Challenge.
+
+
+For the original tensorflow implementation, check this [repo](https://github.com/microsoft/Deep3DFaceReconstruction).
+
+This implementation is written by S. Xu.
+
+## Performance
+
+### ● Reconstruction accuracy
+
+The pytorch implementation achieves lower shape reconstruction error (9% improvement) compare to the [original tensorflow implementation](https://github.com/microsoft/Deep3DFaceReconstruction). Quantitative evaluation (average shape errors in mm) on several benchmarks is as follows:
+
+|Method|FaceWareHouse|MICC Florence | NoW Challenge |
+|:----:|:-----------:|:-----------:|:-----------:|
+|Deep3DFace Tensorflow | 1.81±0.50 | 1.67±0.50 | 1.54±1.29 |
+|**Deep3DFace PyTorch** |**1.64±0.50**|**1.53±0.45**| **1.41±1.21** |
+
+The comparison result with state-of-the-art public 3D face reconstruction methods on the NoW face benchmark is as follows:
+|Rank|Method|Median(mm) | Mean(mm) | Std(mm) |
+|:----:|:-----------:|:-----------:|:-----------:|:-----------:|
+| 1. | [DECA\[Feng et al., SIGGRAPH 2021\]](https://github.com/YadiraF/DECA)|1.09|1.38|1.18|
+| **2.** | **Deep3DFace PyTorch**|**1.11**|**1.41**|**1.21**|
+| 3. | [RingNet [Sanyal et al., CVPR 2019]](https://github.com/soubhiksanyal/RingNet) | 1.21 | 1.53 | 1.31 |
+| 4. | [Deep3DFace [Deng et al., CVPRW 2019]](https://github.com/microsoft/Deep3DFaceReconstruction) | 1.23 | 1.54 | 1.29 |
+| 5. | [3DDFA-V2 [Guo et al., ECCV 2020]](https://github.com/cleardusk/3DDFA_V2) | 1.23 | 1.57 | 1.39 |
+| 6. | [MGCNet [Shang et al., ECCV 2020]](https://github.com/jiaxiangshang/MGCNet) | 1.31 | 1.87 | 2.63 |
+| 7. | [PRNet [Feng et al., ECCV 2018]](https://github.com/YadiraF/PRNet) | 1.50 | 1.98 | 1.88 |
+| 8. | [3DMM-CNN [Tran et al., CVPR 2017]](https://github.com/anhttran/3dmm_cnn) | 1.84 | 2.33 | 2.05 |
+
+For more details about the evaluation, check [Now Challenge](https://ringnet.is.tue.mpg.de/challenge.html) website.
+
+**_A recent benchmark [REALY](https://www.realy3dface.com/) indicates that our method still has the SOTA performance! You can check their paper and website for more details._**
+
+### ● Visual quality
+The pytorch implementation achieves better visual consistency with the input images compare to the original tensorflow version.
+
+
+
+
+
+### ● Speed
+The training speed is on par with the original tensorflow implementation. For more information, see [here](https://github.com/sicxu/Deep3DFaceRecon_pytorch#train-the-face-reconstruction-network).
+
+## Major changes
+
+### ● Differentiable renderer
+
+We use [Nvdiffrast](https://nvlabs.github.io/nvdiffrast/) which is a pytorch library that provides high-performance primitive operations for rasterization-based differentiable rendering. The original tensorflow implementation used [tf_mesh_renderer](https://github.com/google/tf_mesh_renderer) instead.
+
+### ● Face recognition model
+
+We use [Arcface](https://github.com/deepinsight/insightface/tree/master/recognition/arcface_torch), a state-of-the-art face recognition model, for perceptual loss computation. By contrast, the original tensorflow implementation used [Facenet](https://github.com/davidsandberg/facenet).
+
+### ● Training configuration
+
+Data augmentation is used in the training process which contains random image shifting, scaling, rotation, and flipping. We also enlarge the training batchsize from 5 to 32 to stablize the training process.
+
+### ● Training data
+
+We use an extra high quality face image dataset [FFHQ](https://github.com/NVlabs/ffhq-dataset) to increase the diversity of training data.
+
+## Requirements
+**This implementation is only tested under Ubuntu environment with Nvidia GPUs and CUDA installed.**
+
+## Installation
+1. Clone the repository and set up a conda environment with all dependencies as follows:
+```
+git clone https://github.com/sicxu/Deep3DFaceRecon_pytorch.git
+cd Deep3DFaceRecon_pytorch
+conda env create -f environment.yml
+source activate deep3d_pytorch
+```
+
+2. Install Nvdiffrast library:
+```
+git clone https://github.com/NVlabs/nvdiffrast
+cd nvdiffrast # ./Deep3DFaceRecon_pytorch/nvdiffrast
+pip install .
+```
+
+3. Install Arcface Pytorch:
+```
+cd .. # ./Deep3DFaceRecon_pytorch
+git clone https://github.com/deepinsight/insightface.git
+cp -r ./insightface/recognition/arcface_torch ./models/
+```
+## Inference with a pre-trained model
+
+### Prepare prerequisite models
+1. Our method uses [Basel Face Model 2009 (BFM09)](https://faces.dmi.unibas.ch/bfm/main.php?nav=1-0&id=basel_face_model) to represent 3d faces. Get access to BFM09 using this [link](https://faces.dmi.unibas.ch/bfm/main.php?nav=1-2&id=downloads). After getting the access, download "01_MorphableModel.mat". In addition, we use an Expression Basis provided by [Guo et al.](https://github.com/Juyong/3DFace). Download the Expression Basis (Exp_Pca.bin) using this [link (google drive)](https://drive.google.com/file/d/1bw5Xf8C12pWmcMhNEu6PtsYVZkVucEN6/view?usp=sharing). Organize all files into the following structure:
+```
+Deep3DFaceRecon_pytorch
+│
+└─── BFM
+ │
+ └─── 01_MorphableModel.mat
+ │
+ └─── Exp_Pca.bin
+ |
+ └─── ...
+```
+2. We provide a model trained on a combination of [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html),
+[LFW](http://vis-www.cs.umass.edu/lfw/), [300WLP](http://www.cbsr.ia.ac.cn/users/xiangyuzhu/projects/3DDFA/main.htm),
+[IJB-A](https://www.nist.gov/programs-projects/face-challenges), [LS3D-W](https://www.adrianbulat.com/face-alignment), and [FFHQ](https://github.com/NVlabs/ffhq-dataset) datasets. Download the pre-trained model using this [link (google drive)](https://drive.google.com/drive/folders/1liaIxn9smpudjjqMaWWRpP0mXRW_qRPP?usp=sharing) and organize the directory into the following structure:
+```
+Deep3DFaceRecon_pytorch
+│
+└─── checkpoints
+ │
+ └───
+ │
+ └─── epoch_20.pth
+
+```
+
+### Test with custom images
+To reconstruct 3d faces from test images, organize the test image folder as follows:
+```
+Deep3DFaceRecon_pytorch
+│
+└───
+ │
+ └─── *.jpg/*.png
+ |
+ └─── detections
+ |
+ └─── *.txt
+```
+The \*.jpg/\*.png files are test images. The \*.txt files are detected 5 facial landmarks with a shape of 5x2, and have the same name as the corresponding images. Check [./datasets/examples](datasets/examples) for a reference.
+
+Then, run the test script:
+```
+# get reconstruction results of your custom images
+python test.py --name= --epoch=20 --img_folder=
+
+# get reconstruction results of example images
+python test.py --name= --epoch=20 --img_folder=./datasets/examples
+```
+**_Following [#108](https://github.com/sicxu/Deep3DFaceRecon_pytorch/issues/108), if you don't have OpenGL environment, you can simply add "--use_opengl False" to use CUDA context. Make sure you have updated the nvdiffrast to the latest version._**
+
+Results will be saved into ./checkpoints//results/, which contain the following files:
+| \*.png | A combination of cropped input image, reconstructed image, and visualization of projected landmarks.
+|:----|:-----------|
+| \*.obj | Reconstructed 3d face mesh with predicted color (texture+illumination) in the world coordinate space. Best viewed in Meshlab. |
+| \*.mat | Predicted 257-dimensional coefficients and 68 projected 2d facial landmarks. Best viewed in Matlab.
+
+## Training a model from scratch
+### Prepare prerequisite models
+1. We rely on [Arcface](https://github.com/deepinsight/insightface/tree/master/recognition/arcface_torch) to extract identity features for loss computation. Download the pre-trained model from Arcface using this [link](https://github.com/deepinsight/insightface/tree/master/recognition/arcface_torch#ms1mv3). By default, we use the resnet50 backbone ([ms1mv3_arcface_r50_fp16](https://onedrive.live.com/?authkey=%21AFZjr283nwZHqbA&id=4A83B6B633B029CC%215583&cid=4A83B6B633B029CC)), organize the download files into the following structure:
+```
+Deep3DFaceRecon_pytorch
+│
+└─── checkpoints
+ │
+ └─── recog_model
+ │
+ └─── ms1mv3_arcface_r50_fp16
+ |
+ └─── backbone.pth
+```
+2. We initialize R-Net using the weights trained on [ImageNet](https://image-net.org/). Download the weights provided by PyTorch using this [link](https://download.pytorch.org/models/resnet50-0676ba61.pth), and organize the file as the following structure:
+```
+Deep3DFaceRecon_pytorch
+│
+└─── checkpoints
+ │
+ └─── init_model
+ │
+ └─── resnet50-0676ba61.pth
+```
+3. We provide a landmark detector (tensorflow model) to extract 68 facial landmarks for loss computation. The detector is trained on [300WLP](http://www.cbsr.ia.ac.cn/users/xiangyuzhu/projects/3DDFA/main.htm), [LFW](http://vis-www.cs.umass.edu/lfw/), and [LS3D-W](https://www.adrianbulat.com/face-alignment) datasets. Download the trained model using this [link (google drive)](https://drive.google.com/file/d/1Jl1yy2v7lIJLTRVIpgg2wvxYITI8Dkmw/view?usp=sharing) and organize the file as follows:
+```
+Deep3DFaceRecon_pytorch
+│
+└─── checkpoints
+ │
+ └─── lm_model
+ │
+ └─── 68lm_detector.pb
+```
+### Data preparation
+1. To train a model with custom images，5 facial landmarks of each image are needed in advance for an image pre-alignment process. We recommend using [dlib](http://dlib.net/) or [MTCNN](https://github.com/ipazc/mtcnn) to detect these landmarks. Then, organize all files into the following structure:
+```
+Deep3DFaceRecon_pytorch
+│
+└─── datasets
+ │
+ └───
+ │
+ └─── *.png/*.jpg
+ |
+ └─── detections
+ |
+ └─── *.txt
+```
+The \*.txt files contain 5 facial landmarks with a shape of 5x2, and should have the same name with their corresponding images.
+
+2. Generate 68 landmarks and skin attention mask for images using the following script:
+```
+# preprocess training images
+python data_preparation.py --img_folder
+
+# alternatively, you can preprocess multiple image folders simultaneously
+python data_preparation.py --img_folder
+
+# preprocess validation images
+python data_preparation.py --img_folder --mode=val
+```
+The script will generate files of landmarks and skin masks, and save them into ./datasets/. In addition, it also generates a file containing the path of all training data into ./datalist which will then be used in the training script.
+
+### Train the face reconstruction network
+Run the following script to train a face reconstruction model using the pre-processed data:
+```
+# train with single GPU
+python train.py --name= --gpu_ids=0
+
+# train with multiple GPUs
+python train.py --name= --gpu_ids=0,1
+
+# train with other custom settings
+python train.py --name= --gpu_ids=0 --batch_size=32 --n_epochs=20
+```
+Training logs and model parameters will be saved into ./checkpoints/.
+
+By default, the script uses a batchsize of 32 and will train the model with 20 epochs. For reference, the pre-trained model in this repo is trained with the default setting on a image collection of 300k images. A single iteration takes 0.8~0.9s on a single Tesla M40 GPU. The total training process takes around two days.
+
+To use a trained model, see [Inference](https://github.com/sicxu/Deep3DFaceRecon_pytorch#inference-with-a-pre-trained-model) section.
+## Contact
+If you have any questions, please contact the paper authors.
+
+## Citation
+
+Please cite the following paper if this model helps your research:
+
+ @inproceedings{deng2019accurate,
+ title={Accurate 3D Face Reconstruction with Weakly-Supervised Learning: From Single Image to Image Set},
+ author={Yu Deng and Jiaolong Yang and Sicheng Xu and Dong Chen and Yunde Jia and Xin Tong},
+ booktitle={IEEE Computer Vision and Pattern Recognition Workshops},
+ year={2019}
+ }
+##
+The face images on this page are from the public [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) dataset released by MMLab, CUHK.
+
+Part of the code in this implementation takes [CUT](https://github.com/taesungp/contrastive-unpaired-translation) as a reference.
+
diff --git a/Deep3DFaceRecon_pytorch/data/__init__.py b/Deep3DFaceRecon_pytorch/data/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0aaef54bd893314aa8a7b97af8625659cd8d3bfe
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/data/__init__.py
@@ -0,0 +1,118 @@
+"""This package includes all the modules related to data loading and preprocessing
+
+ To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.
+ You need to implement four functions:
+ -- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt).
+ -- <__len__>: return the size of dataset.
+ -- <__getitem__>: get a data point from data loader.
+ -- : (optionally) add dataset-specific options and set default options.
+
+Now you can use the dataset class by specifying flag '--dataset_mode dummy'.
+See our template dataset class 'template_dataset.py' for more details.
+"""
+import importlib
+
+import numpy as np
+import torch.utils.data
+from data.base_dataset import BaseDataset
+
+
+def find_dataset_using_name(dataset_name):
+ """Import the module "data/[dataset_name]_dataset.py".
+
+ In the file, the class called DatasetNameDataset() will
+ be instantiated. It has to be a subclass of BaseDataset,
+ and it is case-insensitive.
+ """
+ dataset_filename = "data." + dataset_name + "_dataset"
+ datasetlib = importlib.import_module(dataset_filename)
+
+ dataset = None
+ target_dataset_name = dataset_name.replace("_", "") + "dataset"
+ for name, cls in datasetlib.__dict__.items():
+ if name.lower() == target_dataset_name.lower() and issubclass(cls, BaseDataset):
+ dataset = cls
+
+ if dataset is None:
+ raise NotImplementedError(
+ "In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase."
+ % (dataset_filename, target_dataset_name)
+ )
+
+ return dataset
+
+
+def get_option_setter(dataset_name):
+ """Return the static method of the dataset class."""
+ dataset_class = find_dataset_using_name(dataset_name)
+ return dataset_class.modify_commandline_options
+
+
+def create_dataset(opt, rank=0):
+ """Create a dataset given the option.
+
+ This function wraps the class CustomDatasetDataLoader.
+ This is the main interface between this package and 'train.py'/'test.py'
+
+ Example:
+ >>> from data import create_dataset
+ >>> dataset = create_dataset(opt)
+ """
+ data_loader = CustomDatasetDataLoader(opt, rank=rank)
+ dataset = data_loader.load_data()
+ return dataset
+
+
+class CustomDatasetDataLoader:
+ """Wrapper class of Dataset class that performs multi-threaded data loading"""
+
+ def __init__(self, opt, rank=0):
+ """Initialize this class
+
+ Step 1: create a dataset instance given the name [dataset_mode]
+ Step 2: create a multi-threaded data loader.
+ """
+ self.opt = opt
+ dataset_class = find_dataset_using_name(opt.dataset_mode)
+ self.dataset = dataset_class(opt)
+ self.sampler = None
+ print("rank %d %s dataset [%s] was created" % (rank, self.dataset.name, type(self.dataset).__name__))
+ if opt.use_ddp and opt.isTrain:
+ world_size = opt.world_size
+ self.sampler = torch.utils.data.distributed.DistributedSampler(
+ self.dataset, num_replicas=world_size, rank=rank, shuffle=not opt.serial_batches
+ )
+ self.dataloader = torch.utils.data.DataLoader(
+ self.dataset,
+ sampler=self.sampler,
+ num_workers=int(opt.num_threads / world_size),
+ batch_size=int(opt.batch_size / world_size),
+ drop_last=True,
+ )
+ else:
+ self.dataloader = torch.utils.data.DataLoader(
+ self.dataset,
+ batch_size=opt.batch_size,
+ shuffle=(not opt.serial_batches) and opt.isTrain,
+ num_workers=int(opt.num_threads),
+ drop_last=True,
+ )
+
+ def set_epoch(self, epoch):
+ self.dataset.current_epoch = epoch
+ if self.sampler is not None:
+ self.sampler.set_epoch(epoch)
+
+ def load_data(self):
+ return self
+
+ def __len__(self):
+ """Return the number of data in the dataset"""
+ return min(len(self.dataset), self.opt.max_dataset_size)
+
+ def __iter__(self):
+ """Return a batch of data"""
+ for i, data in enumerate(self.dataloader):
+ if i * self.opt.batch_size >= self.opt.max_dataset_size:
+ break
+ yield data
diff --git a/Deep3DFaceRecon_pytorch/data/base_dataset.py b/Deep3DFaceRecon_pytorch/data/base_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7d86153c9b17f7365b6480bb95ec272ff7fc095
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/data/base_dataset.py
@@ -0,0 +1,132 @@
+"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
+
+It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
+"""
+import random
+from abc import ABC
+from abc import abstractmethod
+
+import numpy as np
+import torch.utils.data as data
+import torchvision.transforms as transforms
+from PIL import Image
+
+
+class BaseDataset(data.Dataset, ABC):
+ """This class is an abstract base class (ABC) for datasets.
+
+ To create a subclass, you need to implement the following four functions:
+ -- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt).
+ -- <__len__>: return the size of dataset.
+ -- <__getitem__>: get a data point.
+ -- : (optionally) add dataset-specific options and set default options.
+ """
+
+ def __init__(self, opt):
+ """Initialize the class; save the options in the class
+
+ Parameters:
+ opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+ """
+ self.opt = opt
+ # self.root = opt.dataroot
+ self.current_epoch = 0
+
+ @staticmethod
+ def modify_commandline_options(parser, is_train):
+ """Add new dataset-specific options, and rewrite default values for existing options.
+
+ Parameters:
+ parser -- original option parser
+ is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+ Returns:
+ the modified parser.
+ """
+ return parser
+
+ @abstractmethod
+ def __len__(self):
+ """Return the total number of images in the dataset."""
+ return 0
+
+ @abstractmethod
+ def __getitem__(self, index):
+ """Return a data point and its metadata information.
+
+ Parameters:
+ index - - a random integer for data indexing
+
+ Returns:
+ a dictionary of data with their names. It ususally contains the data itself and its metadata information.
+ """
+ pass
+
+
+def get_transform(grayscale=False):
+ transform_list = []
+ if grayscale:
+ transform_list.append(transforms.Grayscale(1))
+ transform_list += [transforms.ToTensor()]
+ return transforms.Compose(transform_list)
+
+
+def get_affine_mat(opt, size):
+ shift_x, shift_y, scale, rot_angle, flip = 0.0, 0.0, 1.0, 0.0, False
+ w, h = size
+
+ if "shift" in opt.preprocess:
+ shift_pixs = int(opt.shift_pixs)
+ shift_x = random.randint(-shift_pixs, shift_pixs)
+ shift_y = random.randint(-shift_pixs, shift_pixs)
+ if "scale" in opt.preprocess:
+ scale = 1 + opt.scale_delta * (2 * random.random() - 1)
+ if "rot" in opt.preprocess:
+ rot_angle = opt.rot_angle * (2 * random.random() - 1)
+ rot_rad = -rot_angle * np.pi / 180
+ if "flip" in opt.preprocess:
+ flip = random.random() > 0.5
+
+ shift_to_origin = np.array([1, 0, -w // 2, 0, 1, -h // 2, 0, 0, 1]).reshape([3, 3])
+ flip_mat = np.array([-1 if flip else 1, 0, 0, 0, 1, 0, 0, 0, 1]).reshape([3, 3])
+ shift_mat = np.array([1, 0, shift_x, 0, 1, shift_y, 0, 0, 1]).reshape([3, 3])
+ rot_mat = np.array([np.cos(rot_rad), np.sin(rot_rad), 0, -np.sin(rot_rad), np.cos(rot_rad), 0, 0, 0, 1]).reshape(
+ [3, 3]
+ )
+ scale_mat = np.array([scale, 0, 0, 0, scale, 0, 0, 0, 1]).reshape([3, 3])
+ shift_to_center = np.array([1, 0, w // 2, 0, 1, h // 2, 0, 0, 1]).reshape([3, 3])
+
+ affine = shift_to_center @ scale_mat @ rot_mat @ shift_mat @ flip_mat @ shift_to_origin
+ affine_inv = np.linalg.inv(affine)
+ return affine, affine_inv, flip
+
+
+def apply_img_affine(img, affine_inv, method=Image.Resampling.BICUBIC):
+ return img.transform(img.size, Image.AFFINE, data=affine_inv.flatten()[:6], resample=Image.Resampling.BICUBIC)
+
+
+def apply_lm_affine(landmark, affine, flip, size):
+ _, h = size
+ lm = landmark.copy()
+ lm[:, 1] = h - 1 - lm[:, 1]
+ lm = np.concatenate((lm, np.ones([lm.shape[0], 1])), -1)
+ lm = lm @ np.transpose(affine)
+ lm[:, :2] = lm[:, :2] / lm[:, 2:]
+ lm = lm[:, :2]
+ lm[:, 1] = h - 1 - lm[:, 1]
+ if flip:
+ lm_ = lm.copy()
+ lm_[:17] = lm[16::-1]
+ lm_[17:22] = lm[26:21:-1]
+ lm_[22:27] = lm[21:16:-1]
+ lm_[31:36] = lm[35:30:-1]
+ lm_[36:40] = lm[45:41:-1]
+ lm_[40:42] = lm[47:45:-1]
+ lm_[42:46] = lm[39:35:-1]
+ lm_[46:48] = lm[41:39:-1]
+ lm_[48:55] = lm[54:47:-1]
+ lm_[55:60] = lm[59:54:-1]
+ lm_[60:65] = lm[64:59:-1]
+ lm_[65:68] = lm[67:64:-1]
+ lm = lm_
+ return lm
diff --git a/Deep3DFaceRecon_pytorch/data/flist_dataset.py b/Deep3DFaceRecon_pytorch/data/flist_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..21265f361b46cbbce126b8541284d1bebf241e08
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/data/flist_dataset.py
@@ -0,0 +1,129 @@
+"""This script defines the custom dataset for Deep3DFaceRecon_pytorch
+"""
+import json
+import os.path
+import pickle
+import random
+
+import numpy as np
+import torch
+import util.util as util
+from data.base_dataset import apply_img_affine
+from data.base_dataset import apply_lm_affine
+from data.base_dataset import BaseDataset
+from data.base_dataset import get_affine_mat
+from data.base_dataset import get_transform
+from data.image_folder import make_dataset
+from PIL import Image
+from scipy.io import loadmat
+from scipy.io import savemat
+from util.load_mats import load_lm3d
+from util.preprocess import align_img
+from util.preprocess import estimate_norm
+
+
+def default_flist_reader(flist):
+ """
+ flist format: impath label\nimpath label\n ...(same to caffe's filelist)
+ """
+ imlist = []
+ with open(flist, "r") as rf:
+ for line in rf.readlines():
+ impath = line.strip()
+ imlist.append(impath)
+
+ return imlist
+
+
+def jason_flist_reader(flist):
+ with open(flist, "r") as fp:
+ info = json.load(fp)
+ return info
+
+
+def parse_label(label):
+ return torch.tensor(np.array(label).astype(np.float32))
+
+
+class FlistDataset(BaseDataset):
+ """
+ It requires one directories to host training images '/path/to/data/train'
+ You can train the model with the dataset flag '--dataroot /path/to/data'.
+ """
+
+ def __init__(self, opt):
+ """Initialize this dataset class.
+
+ Parameters:
+ opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+ """
+ BaseDataset.__init__(self, opt)
+
+ self.lm3d_std = load_lm3d(opt.bfm_folder)
+
+ msk_names = default_flist_reader(opt.flist)
+ self.msk_paths = [os.path.join(opt.data_root, i) for i in msk_names]
+
+ self.size = len(self.msk_paths)
+ self.opt = opt
+
+ self.name = "train" if opt.isTrain else "val"
+ if "_" in opt.flist:
+ self.name += "_" + opt.flist.split(os.sep)[-1].split("_")[0]
+
+ def __getitem__(self, index):
+ """Return a data point and its metadata information.
+
+ Parameters:
+ index (int) -- a random integer for data indexing
+
+ Returns a dictionary that contains A, B, A_paths and B_paths
+ img (tensor) -- an image in the input domain
+ msk (tensor) -- its corresponding attention mask
+ lm (tensor) -- its corresponding 3d landmarks
+ im_paths (str) -- image paths
+ aug_flag (bool) -- a flag used to tell whether its raw or augmented
+ """
+ msk_path = self.msk_paths[index % self.size] # make sure index is within then range
+ img_path = msk_path.replace("mask/", "")
+ lm_path = ".".join(msk_path.replace("mask", "landmarks").split(".")[:-1]) + ".txt"
+
+ raw_img = Image.open(img_path).convert("RGB")
+ raw_msk = Image.open(msk_path).convert("RGB")
+ raw_lm = np.loadtxt(lm_path).astype(np.float32)
+
+ _, img, lm, msk = align_img(raw_img, raw_lm, self.lm3d_std, raw_msk)
+
+ aug_flag = self.opt.use_aug and self.opt.isTrain
+ if aug_flag:
+ img, lm, msk = self._augmentation(img, lm, self.opt, msk)
+
+ _, H = img.size
+ M = estimate_norm(lm, H)
+ transform = get_transform()
+ img_tensor = transform(img)
+ msk_tensor = transform(msk)[:1, ...]
+ lm_tensor = parse_label(lm)
+ M_tensor = parse_label(M)
+
+ return {
+ "imgs": img_tensor,
+ "lms": lm_tensor,
+ "msks": msk_tensor,
+ "M": M_tensor,
+ "im_paths": img_path,
+ "aug_flag": aug_flag,
+ "dataset": self.name,
+ }
+
+ def _augmentation(self, img, lm, opt, msk=None):
+ affine, affine_inv, flip = get_affine_mat(opt, img.size)
+ img = apply_img_affine(img, affine_inv)
+ lm = apply_lm_affine(lm, affine, flip, img.size)
+ if msk is not None:
+ msk = apply_img_affine(msk, affine_inv, method=Image.BILINEAR)
+ return img, lm, msk
+
+ def __len__(self):
+ """Return the total number of images in the dataset."""
+ return self.size
diff --git a/Deep3DFaceRecon_pytorch/data/image_folder.py b/Deep3DFaceRecon_pytorch/data/image_folder.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a02dfe06948ddb90a76667d538f7e3ea47a33e8
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/data/image_folder.py
@@ -0,0 +1,77 @@
+"""A modified image folder class
+
+We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
+so that this class can load images from both current directory and its subdirectories.
+"""
+import os.path
+
+import numpy as np
+import torch.utils.data as data
+from PIL import Image
+
+IMG_EXTENSIONS = [
+ ".jpg",
+ ".JPG",
+ ".jpeg",
+ ".JPEG",
+ ".png",
+ ".PNG",
+ ".ppm",
+ ".PPM",
+ ".bmp",
+ ".BMP",
+ ".tif",
+ ".TIF",
+ ".tiff",
+ ".TIFF",
+]
+
+
+def is_image_file(filename):
+ return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir, max_dataset_size=float("inf")):
+ images = []
+ assert os.path.isdir(dir) or os.path.islink(dir), "%s is not a valid directory" % dir
+
+ for root, _, fnames in sorted(os.walk(dir, followlinks=True)):
+ for fname in fnames:
+ if is_image_file(fname):
+ path = os.path.join(root, fname)
+ images.append(path)
+ return images[: min(max_dataset_size, len(images))]
+
+
+def default_loader(path):
+ return Image.open(path).convert("RGB")
+
+
+class ImageFolder(data.Dataset):
+ def __init__(self, root, transform=None, return_paths=False, loader=default_loader):
+ imgs = make_dataset(root)
+ if len(imgs) == 0:
+ raise (
+ RuntimeError(
+ "Found 0 images in: " + root + "\n" "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)
+ )
+ )
+
+ self.root = root
+ self.imgs = imgs
+ self.transform = transform
+ self.return_paths = return_paths
+ self.loader = loader
+
+ def __getitem__(self, index):
+ path = self.imgs[index]
+ img = self.loader(path)
+ if self.transform is not None:
+ img = self.transform(img)
+ if self.return_paths:
+ return img, path
+ else:
+ return img
+
+ def __len__(self):
+ return len(self.imgs)
diff --git a/Deep3DFaceRecon_pytorch/data/template_dataset.py b/Deep3DFaceRecon_pytorch/data/template_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..d776bde9c36396865e149922305cb2284942529b
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/data/template_dataset.py
@@ -0,0 +1,80 @@
+"""Dataset class template
+
+This module provides a template for users to implement custom datasets.
+You can specify '--dataset_mode template' to use this dataset.
+The class name should be consistent with both the filename and its dataset_mode option.
+The filename should be _dataset.py
+The class name should be Dataset.py
+You need to implement the following functions:
+ -- : Add dataset-specific options and rewrite default values for existing options.
+ -- <__init__>: Initialize this dataset class.
+ -- <__getitem__>: Return a data point and its metadata information.
+ -- <__len__>: Return the number of images.
+"""
+from data.base_dataset import BaseDataset
+from data.base_dataset import get_transform
+
+# from data.image_folder import make_dataset
+# from PIL import Image
+
+
+class TemplateDataset(BaseDataset):
+ """A template dataset class for you to implement custom datasets."""
+
+ @staticmethod
+ def modify_commandline_options(parser, is_train):
+ """Add new dataset-specific options, and rewrite default values for existing options.
+
+ Parameters:
+ parser -- original option parser
+ is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+ Returns:
+ the modified parser.
+ """
+ parser.add_argument("--new_dataset_option", type=float, default=1.0, help="new dataset option")
+ parser.set_defaults(max_dataset_size=10, new_dataset_option=2.0) # specify dataset-specific default values
+ return parser
+
+ def __init__(self, opt):
+ """Initialize this dataset class.
+
+ Parameters:
+ opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+
+ A few things can be done here.
+ - save the options (have been done in BaseDataset)
+ - get image paths and meta information of the dataset.
+ - define the image transformation.
+ """
+ # save the option and dataset root
+ BaseDataset.__init__(self, opt)
+ # get the image paths of your dataset;
+ self.image_paths = (
+ []
+ ) # You can call sorted(make_dataset(self.root, opt.max_dataset_size)) to get all the image paths under the directory self.root
+ # define the default transform function. You can use ; You can also define your custom transform function
+ self.transform = get_transform(opt)
+
+ def __getitem__(self, index):
+ """Return a data point and its metadata information.
+
+ Parameters:
+ index -- a random integer for data indexing
+
+ Returns:
+ a dictionary of data with their names. It usually contains the data itself and its metadata information.
+
+ Step 1: get a random image path: e.g., path = self.image_paths[index]
+ Step 2: load your data from the disk: e.g., image = Image.open(path).convert('RGB').
+ Step 3: convert your data to a PyTorch tensor. You can use helpder functions such as self.transform. e.g., data = self.transform(image)
+ Step 4: return a data point as a dictionary.
+ """
+ path = "temp" # needs to be a string
+ data_A = None # needs to be a tensor
+ data_B = None # needs to be a tensor
+ return {"data_A": data_A, "data_B": data_B, "path": path}
+
+ def __len__(self):
+ """Return the total number of images."""
+ return len(self.image_paths)
diff --git a/Deep3DFaceRecon_pytorch/data_preparation.py b/Deep3DFaceRecon_pytorch/data_preparation.py
new file mode 100644
index 0000000000000000000000000000000000000000..74b9361d45a56342b10530e015b4099b91bbf625
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/data_preparation.py
@@ -0,0 +1,57 @@
+"""This script is the data preparation script for Deep3DFaceRecon_pytorch
+"""
+import argparse
+import os
+import warnings
+
+import numpy as np
+from util.detect_lm68 import detect_68p
+from util.detect_lm68 import load_lm_graph
+from util.generate_list import check_list
+from util.generate_list import write_list
+from util.skin_mask import get_skin_mask
+
+warnings.filterwarnings("ignore")
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--data_root", type=str, default="datasets", help="root directory for training data")
+parser.add_argument("--img_folder", nargs="+", required=True, help="folders of training images")
+parser.add_argument("--mode", type=str, default="train", help="train or val")
+opt = parser.parse_args()
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+
+def data_prepare(folder_list, mode):
+
+ lm_sess, input_op, output_op = load_lm_graph(
+ "./checkpoints/lm_model/68lm_detector.pb"
+ ) # load a tensorflow version 68-landmark detector
+
+ for img_folder in folder_list:
+ detect_68p(img_folder, lm_sess, input_op, output_op) # detect landmarks for images
+ get_skin_mask(img_folder) # generate skin attention mask for images
+
+ # create files that record path to all training data
+ msks_list = []
+ for img_folder in folder_list:
+ path = os.path.join(img_folder, "mask")
+ msks_list += [
+ "/".join([img_folder, "mask", i])
+ for i in sorted(os.listdir(path))
+ if "jpg" in i or "png" in i or "jpeg" in i or "PNG" in i
+ ]
+
+ imgs_list = [i.replace("mask/", "") for i in msks_list]
+ lms_list = [i.replace("mask", "landmarks") for i in msks_list]
+ lms_list = [".".join(i.split(".")[:-1]) + ".txt" for i in lms_list]
+
+ lms_list_final, imgs_list_final, msks_list_final = check_list(
+ lms_list, imgs_list, msks_list
+ ) # check if the path is valid
+ write_list(lms_list_final, imgs_list_final, msks_list_final, mode=mode) # save files
+
+
+if __name__ == "__main__":
+ print("Datasets:", opt.img_folder)
+ data_prepare([os.path.join(opt.data_root, folder) for folder in opt.img_folder], opt.mode)
diff --git a/Deep3DFaceRecon_pytorch/environment.yml b/Deep3DFaceRecon_pytorch/environment.yml
new file mode 100644
index 0000000000000000000000000000000000000000..c16d73cfd9d8e5488362aee9067ca5ef9897f60d
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/environment.yml
@@ -0,0 +1,24 @@
+name: deep3d_pytorch
+channels:
+ - pytorch
+ - conda-forge
+ - defaults
+dependencies:
+ - python=3.6
+ - pytorch=1.6.0
+ - torchvision=0.7.0
+ - numpy=1.18.1
+ - scikit-image=0.16.2
+ - scipy=1.4.1
+ - pillow=6.2.1
+ - pip=20.0.2
+ - ipython=7.13.0
+ - yaml=0.1.7
+ - pip:
+ - matplotlib==2.2.5
+ - opencv-python==3.4.9.33
+ - tensorboard==1.15.0
+ - tensorflow==1.15.0
+ - kornia==0.5.5
+ - dominate==2.6.0
+ - trimesh==3.9.20
\ No newline at end of file
diff --git a/Deep3DFaceRecon_pytorch/models/__init__.py b/Deep3DFaceRecon_pytorch/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4642bcbf549aa1354268ed7b6567d20eff6b1237
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/models/__init__.py
@@ -0,0 +1,69 @@
+"""This package contains modules related to objective functions, optimizations, and network architectures.
+
+To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
+You need to implement the following five functions:
+ -- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
+ -- : unpack data from dataset and apply preprocessing.
+ -- : produce intermediate results.
+ -- : calculate loss, gradients, and update network weights.
+ -- : (optionally) add model-specific options and set default options.
+
+In the function <__init__>, you need to define four lists:
+ -- self.loss_names (str list): specify the training losses that you want to plot and save.
+ -- self.model_names (str list): define networks used in our training.
+ -- self.visual_names (str list): specify the images that you want to display and save.
+ -- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
+
+Now you can use the model class by specifying flag '--model dummy'.
+See our template model class 'template_model.py' for more details.
+"""
+import importlib
+
+from Deep3DFaceRecon_pytorch.models.base_model import BaseModel
+
+
+def find_model_using_name(model_name):
+ """Import the module "models/[model_name]_model.py".
+
+ In the file, the class called DatasetNameModel() will
+ be instantiated. It has to be a subclass of BaseModel,
+ and it is case-insensitive.
+ """
+ model_filename = "models." + model_name + "_model"
+ modellib = importlib.import_module(model_filename)
+ model = None
+ target_model_name = model_name.replace("_", "") + "model"
+ for name, cls in modellib.__dict__.items():
+ if name.lower() == target_model_name.lower() and issubclass(cls, BaseModel):
+ model = cls
+
+ if model is None:
+ print(
+ "In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase."
+ % (model_filename, target_model_name)
+ )
+ exit(0)
+
+ return model
+
+
+def get_option_setter(model_name):
+ """Return the static method of the model class."""
+ model_class = find_model_using_name(model_name)
+ return model_class.modify_commandline_options
+
+
+def create_model(opt):
+ """Create a model given the option.
+
+ This function warps the class CustomDatasetDataLoader.
+ This is the main interface between this package and 'train.py'/'test.py'
+
+ Example:
+ >>> from models import create_model
+ >>> model = create_model(opt)
+ """
+ model = find_model_using_name(opt.model)
+ instance = model(opt)
+ print("model [%s] was created" % type(instance).__name__)
+ return instance
diff --git a/Deep3DFaceRecon_pytorch/models/__pycache__/__init__.cpython-310.pyc b/Deep3DFaceRecon_pytorch/models/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e0f5dd8cda00a950e69b90f7d3c8d0abb21858aa
Binary files /dev/null and b/Deep3DFaceRecon_pytorch/models/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Deep3DFaceRecon_pytorch/models/__pycache__/base_model.cpython-310.pyc b/Deep3DFaceRecon_pytorch/models/__pycache__/base_model.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ee4ae4c8a8e3e14c47e405a901abd4769c61c076
Binary files /dev/null and b/Deep3DFaceRecon_pytorch/models/__pycache__/base_model.cpython-310.pyc differ
diff --git a/Deep3DFaceRecon_pytorch/models/__pycache__/bfm.cpython-310.pyc b/Deep3DFaceRecon_pytorch/models/__pycache__/bfm.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f4d7594423e1484031c01ed95e86c8a5f4c463ef
Binary files /dev/null and b/Deep3DFaceRecon_pytorch/models/__pycache__/bfm.cpython-310.pyc differ
diff --git a/Deep3DFaceRecon_pytorch/models/__pycache__/networks.cpython-310.pyc b/Deep3DFaceRecon_pytorch/models/__pycache__/networks.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4bdbdad04a18832a3ad1c92ca85a3a69698465c5
Binary files /dev/null and b/Deep3DFaceRecon_pytorch/models/__pycache__/networks.cpython-310.pyc differ
diff --git a/Deep3DFaceRecon_pytorch/models/arcface_torch/README.md b/Deep3DFaceRecon_pytorch/models/arcface_torch/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..8d391f63684dd1f47900dc6449a5e22fa25e3da3
--- /dev/null
+++ b/Deep3DFaceRecon_pytorch/models/arcface_torch/README.md
@@ -0,0 +1,218 @@
+# Distributed Arcface Training in Pytorch
+
+The "arcface_torch" repository is the official implementation of the ArcFace algorithm. It supports distributed and sparse training with multiple distributed training examples, including several memory-saving techniques such as mixed precision training and gradient checkpointing. It also supports training for ViT models and datasets including WebFace42M and Glint360K, two of the largest open-source datasets. Additionally, the repository comes with a built-in tool for converting to ONNX format, making it easy to submit to MFR evaluation systems.
+
+[](https://paperswithcode.com/sota/face-verification-on-ijb-c?p=killing-two-birds-with-one-stone-efficient)
+[](https://paperswithcode.com/sota/face-verification-on-ijb-b?p=killing-two-birds-with-one-stone-efficient)
+[](https://paperswithcode.com/sota/face-verification-on-agedb-30?p=killing-two-birds-with-one-stone-efficient)
+[](https://paperswithcode.com/sota/face-verification-on-cfp-fp?p=killing-two-birds-with-one-stone-efficient)
+
+## Requirements
+
+To avail the latest features of PyTorch, we have upgraded to version 1.12.0.
+
+- Install [PyTorch](https://pytorch.org/get-started/previous-versions/) (torch>=1.12.0).
+- (Optional) Install [DALI](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/), our doc for [install_dali.md](docs/install_dali.md).
+- `pip install -r requirement.txt`.
+
+## How to Training
+
+To train a model, execute the `train.py` script with the path to the configuration files. The sample commands provided below demonstrate the process of conducting distributed training.
+
+### 1. To run on one GPU:
+
+```shell
+python train_v2.py configs/ms1mv3_r50_onegpu
+```
+
+Note:
+It is not recommended to use a single GPU for training, as this may result in longer training times and suboptimal performance. For best results, we suggest using multiple GPUs or a GPU cluster.
+
+
+### 2. To run on a machine with 8 GPUs:
+
+```shell
+torchrun --nproc_per_node=8 train.py configs/ms1mv3_r50
+```
+
+### 3. To run on 2 machines with 8 GPUs each:
+
+Node 0:
+
+```shell
+torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr="ip1" --master_port=12581 train.py configs/wf42m_pfc02_16gpus_r100
+```
+
+Node 1:
+
+```shell
+torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr="ip1" --master_port=12581 train.py configs/wf42m_pfc02_16gpus_r100
+```
+
+### 4. Run ViT-B on a machine with 24k batchsize:
+
+```shell
+torchrun --nproc_per_node=8 train_v2.py configs/wf42m_pfc03_40epoch_8gpu_vit_b
+```
+
+
+## Download Datasets or Prepare Datasets
+- [MS1MV2](https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_#ms1m-arcface-85k-ids58m-images-57) (87k IDs, 5.8M images)
+- [MS1MV3](https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_#ms1m-retinaface) (93k IDs, 5.2M images)
+- [Glint360K](https://github.com/deepinsight/insightface/tree/master/recognition/partial_fc#4-download) (360k IDs, 17.1M images)
+- [WebFace42M](docs/prepare_webface42m.md) (2M IDs, 42.5M images)
+- [Your Dataset, Click Here!](docs/prepare_custom_dataset.md)
+
+Note:
+If you want to use DALI for data reading, please use the script 'scripts/shuffle_rec.py' to shuffle the InsightFace style rec before using it.
+Example:
+
+`python scripts/shuffle_rec.py ms1m-retinaface-t1`
+
+You will get the "shuffled_ms1m-retinaface-t1" folder, where the samples in the "train.rec" file are shuffled.
+
+
+## Model Zoo
+
+- The models are available for non-commercial research purposes only.
+- All models can be found in here.
+- [Baidu Yun Pan](https://pan.baidu.com/s/1CL-l4zWqsI1oDuEEYVhj-g): e8pw
+- [OneDrive](https://1drv.ms/u/s!AswpsDO2toNKq0lWY69vN58GR6mw?e=p9Ov5d)
+
+### Performance on IJB-C and [**ICCV2021-MFR**](https://github.com/deepinsight/insightface/blob/master/challenges/mfr/README.md)
+
+ICCV2021-MFR testset consists of non-celebrities so we can ensure that it has very few overlap with public available face
+recognition training set, such as MS1M and CASIA as they mostly collected from online celebrities.
+As the result, we can evaluate the FAIR performance for different algorithms.
+
+For **ICCV2021-MFR-ALL** set, TAR is measured on all-to-all 1:1 protocal, with FAR less than 0.000001(e-6). The
+globalised multi-racial testset contains 242,143 identities and 1,624,305 images.
+
+
+#### 1. Training on Single-Host GPU
+
+| Datasets | Backbone | **MFR-ALL** | IJB-C(1E-4) | IJB-C(1E-5) | log |
+|:---------------|:--------------------|:------------|:------------|:------------|:------------------------------------------------------------------------------------------------------------------------------------|
+| MS1MV2 | mobilefacenet-0.45G | 62.07 | 93.61 | 90.28 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv2_mbf/training.log) |
+| MS1MV2 | r50 | 75.13 | 95.97 | 94.07 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv2_r50/training.log) |
+| MS1MV2 | r100 | 78.12 | 96.37 | 94.27 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv2_r100/training.log) |
+| MS1MV3 | mobilefacenet-0.45G | 63.78 | 94.23 | 91.33 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_mbf/training.log) |
+| MS1MV3 | r50 | 79.14 | 96.37 | 94.47 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_r50/training.log) |
+| MS1MV3 | r100 | 81.97 | 96.85 | 95.02 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_r100/training.log) |
+| Glint360K | mobilefacenet-0.45G | 70.18 | 95.04 | 92.62 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_mbf/training.log) |
+| Glint360K | r50 | 86.34 | 97.16 | 95.81 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_r50/training.log) |
+| Glint360k | r100 | 89.52 | 97.55 | 96.38 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_r100/training.log) |
+| WF4M | r100 | 89.87 | 97.19 | 95.48 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/wf4m_r100/training.log) |
+| WF12M-PFC-0.2 | r100 | 94.75 | 97.60 | 95.90 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/wf12m_pfc02_r100/training.log) |
+| WF12M-PFC-0.3 | r100 | 94.71 | 97.64 | 96.01 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/wf12m_pfc03_r100/training.log) |
+| WF12M | r100 | 94.69 | 97.59 | 95.97 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/wf12m_r100/training.log) |
+| WF42M-PFC-0.2 | r100 | 96.27 | 97.70 | 96.31 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/wf42m_pfc02_r100/training.log) |
+| WF42M-PFC-0.2 | ViT-T-1.5G | 92.04 | 97.27 | 95.68 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/wf42m_pfc02_40epoch_8gpu_vit_t/training.log) |
+| WF42M-PFC-0.3 | ViT-B-11G | 97.16 | 97.91 | 97.05 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/pfc03_wf42m_vit_b_8gpu/training.log) |
+
+#### 2. Training on Multi-Host GPU
+
+| Datasets | Backbone(bs*gpus) | **MFR-ALL** | IJB-C(1E-4) | IJB-C(1E-5) | Throughout | log |
+|:-----------------|:------------------|:------------|:------------|:------------|:-----------|:-------------------------------------------------------------------------------------------------------------------------------------------|
+| WF42M-PFC-0.2 | r50(512*8) | 93.83 | 97.53 | 96.16 | ~5900 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/webface42m_r50_bs4k_pfc02/training.log) |
+| WF42M-PFC-0.2 | r50(512*16) | 93.96 | 97.46 | 96.12 | ~11000 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/webface42m_r50_lr01_pfc02_bs8k_16gpus/training.log) |
+| WF42M-PFC-0.2 | r50(128*32) | 94.04 | 97.48 | 95.94 | ~17000 | click me |
+| WF42M-PFC-0.2 | r100(128*16) | 96.28 | 97.80 | 96.57 | ~5200 | click me |
+| WF42M-PFC-0.2 | r100(256*16) | 96.69 | 97.85 | 96.63 | ~5200 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/webface42m_r100_bs4k_pfc02/training.log) |
+| WF42M-PFC-0.0018 | r100(512*32) | 93.08 | 97.51 | 95.88 | ~10000 | click me |
+| WF42M-PFC-0.2 | r100(128*32) | 96.57 | 97.83 | 96.50 | ~9800 | click me |
+
+`r100(128*32)` means backbone is r100, batchsize per gpu is 128, the number of gpus is 32.
+
+
+
+#### 3. ViT For Face Recognition
+
+| Datasets | Backbone(bs) | FLOPs | **MFR-ALL** | IJB-C(1E-4) | IJB-C(1E-5) | Throughout | log |
+|:--------------|:--------------|:------|:------------|:------------|:------------|:-----------|:-----------------------------------------------------------------------------------------------------------------------------|
+| WF42M-PFC-0.3 | r18(128*32) | 2.6 | 79.13 | 95.77 | 93.36 | - | click me |
+| WF42M-PFC-0.3 | r50(128*32) | 6.3 | 94.03 | 97.48 | 95.94 | - | click me |
+| WF42M-PFC-0.3 | r100(128*32) | 12.1 | 96.69 | 97.82 | 96.45 | - | click me |
+| WF42M-PFC-0.3 | r200(128*32) | 23.5 | 97.70 | 97.97 | 96.93 | - | click me |
+| WF42M-PFC-0.3 | VIT-T(384*64) | 1.5 | 92.24 | 97.31 | 95.97 | ~35000 | click me |
+| WF42M-PFC-0.3 | VIT-S(384*64) | 5.7 | 95.87 | 97.73 | 96.57 | ~25000 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/pfc03_wf42m_vit_s_64gpu/training.log) |
+| WF42M-PFC-0.3 | VIT-B(384*64) | 11.4 | 97.42 | 97.90 | 97.04 | ~13800 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/pfc03_wf42m_vit_b_64gpu/training.log) |
+| WF42M-PFC-0.3 | VIT-L(384*64) | 25.3 | 97.85 | 98.00 | 97.23 | ~9406 | [click me](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/pfc03_wf42m_vit_l_64gpu/training.log) |
+
+`WF42M` means WebFace42M, `PFC-0.3` means negivate class centers sample rate is 0.3.
+
+#### 4. Noisy Datasets
+
+| Datasets | Backbone | **MFR-ALL** | IJB-C(1E-4) | IJB-C(1E-5) | log |
+|:-------------------------|:---------|:------------|:------------|:------------|:---------|
+| WF12M-Flip(40%) | r50 | 43.87 | 88.35 | 80.78 | click me |
+| WF12M-Flip(40%)-PFC-0.1* | r50 | 80.20 | 96.11 | 93.79 | click me |
+| WF12M-Conflict | r50 | 79.93 | 95.30 | 91.56 | click me |
+| WF12M-Conflict-PFC-0.3* | r50 | 91.68 | 97.28 | 95.75 | click me |
+
+`WF12M` means WebFace12M, `+PFC-0.1*` denotes additional abnormal inter-class filtering.
+
+
+
+## Speed Benchmark
+