face3d model training added

dataloader/chong_imutils.py

@@ -0,0 +1,117 @@
+from __future__ import absolute_import
+import torch
+import numpy as np
+import torch
+
+import matplotlib as mpl
+mpl.use('Agg')
+from matplotlib import pyplot as plt
+import numpy as np
+import pandas as pd
+
+
+def smooth_by_conv(window_size, df, col):
+    padded_track = pd.concat([pd.DataFrame([[df.iloc[0][col]]]*(window_size//2), columns=[0]),
+                     df[col],
+                     pd.DataFrame([[df.iloc[-1][col]]]*(window_size//2), columns=[0])])
+    smoothed_signals = np.convolve(padded_track.squeeze(), np.ones(window_size)/window_size, mode='valid')
+    return smoothed_signals
+
+
+def to_numpy(tensor):
+    if torch.is_tensor(tensor):
+        return tensor.cpu().numpy()
+    elif type(tensor).__module__ != 'numpy':
+        raise ValueError("Cannot convert {} to numpy array"
+                         .format(type(tensor)))
+    return tensor
+
+
+def to_torch(ndarray):
+    if type(ndarray).__module__ == 'numpy':
+        return torch.from_numpy(ndarray)
+    elif not torch.is_tensor(ndarray):
+        raise ValueError("Cannot convert {} to torch tensor"
+                         .format(type(ndarray)))
+    return ndarray
+
+
+def unnorm(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
+    std = np.array(std).reshape(3,1,1)
+    mean = np.array(mean).reshape(3,1,1)
+    return img * std + mean
+
+
+def get_head_box_channel(x_min, y_min, x_max, y_max, width, height, resolution, coordconv=False):
+    head_box = np.array([x_min/width, y_min/height, x_max/width, y_max/height])*resolution
+    head_box = head_box.astype(int)
+    head_box = np.clip(head_box, 0, resolution-1)
+    if coordconv:
+        unit = np.array(range(0,resolution), dtype=np.float32)
+        head_channel = []
+        for i in unit:
+            head_channel.append([unit+i])
+        head_channel = np.squeeze(np.array(head_channel)) / float(np.max(head_channel))
+        head_channel[head_box[1]:head_box[3],head_box[0]:head_box[2]] = 0
+    else:
+        head_channel = np.zeros((resolution,resolution), dtype=np.float32)
+        head_channel[head_box[1]:head_box[3],head_box[0]:head_box[2]] = 1
+    head_channel = torch.from_numpy(head_channel)
+    return head_channel
+
+def get_object_box_channel(gtboxes,width,height,resolution):
+    channel = np.zeros((resolution,resolution), dtype=np.float32)
+    for box in gtboxes:
+        obj_box = np.array([box[0], box[1], box[2], box[3]])* resolution
+        obj_box = obj_box.astype(int)
+        # obj_box = np.clip(obj_box, 0, resolution-1)
+        channel[obj_box[1]:obj_box[3],obj_box[0]:obj_box[2]] = 1
+    return torch.from_numpy(channel)    
+
+
+def draw_labelmap(img, pt, sigma, type='Gaussian'):
+    # Draw a 2D gaussian
+    # Adopted from https://github.com/anewell/pose-hg-train/blob/master/src/pypose/draw.py
+    img = to_numpy(img)
+
+    # Check that any part of the gaussian is in-bounds
+    ul = [int(pt[0] - 3 * sigma), int(pt[1] - 3 * sigma)]
+    br = [int(pt[0] + 3 * sigma + 1), int(pt[1] + 3 * sigma + 1)]
+    if (ul[0] >= img.shape[1] or ul[1] >= img.shape[0] or
+            br[0] < 0 or br[1] < 0):
+        # If not, just return the image as is
+        return to_torch(img)
+
+    # Generate gaussian
+    size = 6 * sigma + 1
+    x = np.arange(0, size, 1, float)
+    y = x[:, np.newaxis]
+    x0 = y0 = size // 2
+    # The gaussian is not normalized, we want the center value to equal 1
+    if type == 'Gaussian':
+        g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
+    elif type == 'Cauchy':
+        g = sigma / (((x - x0) ** 2 + (y - y0) ** 2 + sigma ** 2) ** 1.5)
+
+    # Usable gaussian range
+    g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0]
+    g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1]
+    # Image range
+    img_x = max(0, ul[0]), min(br[0], img.shape[1])
+    img_y = max(0, ul[1]), min(br[1], img.shape[0])
+
+    img[img_y[0]:img_y[1], img_x[0]:img_x[1]] += g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
+    img = img/np.max(img) # normalize heatmap so it has max value of 1
+    return to_torch(img)
+
+
+def multi_hot_targets(gaze_pts, out_res):
+    w, h = out_res
+    target_map = np.zeros((h, w))
+    for p in gaze_pts:
+        if p[0] >= 0:
+            x, y = map(int,[p[0]*w.float(), p[1]*h.float()])
+            x = min(x, w-1)
+            y = min(y, h-1)
+            target_map[y, x] = 1
+    return target_map