File size: 4,901 Bytes
9c79341 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
"""This script contains the image preprocessing code for Deep3DFaceRecon_pytorch
"""
import numpy as np
from scipy.io import loadmat
from PIL import Image
import cv2
import os
from skimage import transform as trans
import torch
import warnings
warnings.filterwarnings("ignore", category=np.VisibleDeprecationWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
# calculating least square problem for image alignment
def POS(xp, x):
npts = xp.shape[1]
A = np.zeros([2*npts, 8])
A[0:2*npts-1:2, 0:3] = x.transpose()
A[0:2*npts-1:2, 3] = 1
A[1:2*npts:2, 4:7] = x.transpose()
A[1:2*npts:2, 7] = 1
b = np.reshape(xp.transpose(), [2*npts, 1])
k, _, _, _ = np.linalg.lstsq(A, b)
R1 = k[0:3]
R2 = k[4:7]
sTx = k[3]
sTy = k[7]
s = (np.linalg.norm(R1) + np.linalg.norm(R2))/2
t = np.stack([sTx, sTy], axis=0)
return t, s
# resize and crop images for face reconstruction
def resize_n_crop_img(img, lm, t, s, target_size=224., mask=None):
w0, h0 = img.size
w = (w0*s).astype(np.int32)
h = (h0*s).astype(np.int32)
left = (w/2 - target_size/2 + float((t[0] - w0/2)*s)).astype(np.int32)
right = left + target_size
up = (h/2 - target_size/2 + float((h0/2 - t[1])*s)).astype(np.int32)
below = up + target_size
img = img.resize((w, h), resample=Image.BICUBIC)
img = img.crop((left, up, right, below))
if mask is not None:
mask = mask.resize((w, h), resample=Image.BICUBIC)
mask = mask.crop((left, up, right, below))
lm = np.stack([lm[:, 0] - t[0] + w0/2, lm[:, 1] -
t[1] + h0/2], axis=1)*s
lm = lm - np.reshape(
np.array([(w/2 - target_size/2), (h/2-target_size/2)]), [1, 2])
return img, lm, mask
# utils for face reconstruction
def extract_5p(lm):
lm_idx = np.array([31, 37, 40, 43, 46, 49, 55]) - 1
lm5p = np.stack([lm[lm_idx[0], :], np.mean(lm[lm_idx[[1, 2]], :], 0), np.mean(
lm[lm_idx[[3, 4]], :], 0), lm[lm_idx[5], :], lm[lm_idx[6], :]], axis=0)
lm5p = lm5p[[1, 2, 0, 3, 4], :]
return lm5p
# utils for face reconstruction
def align_img(img, lm, lm3D, mask=None, target_size=224., rescale_factor=102.):
"""
Return:
transparams --numpy.array (raw_W, raw_H, scale, tx, ty)
img_new --PIL.Image (target_size, target_size, 3)
lm_new --numpy.array (68, 2), y direction is opposite to v direction
mask_new --PIL.Image (target_size, target_size)
Parameters:
img --PIL.Image (raw_H, raw_W, 3)
lm --numpy.array (68, 2), y direction is opposite to v direction
lm3D --numpy.array (5, 3)
mask --PIL.Image (raw_H, raw_W, 3)
"""
try:
# Debug input shapes
print(f"\n[DEBUG] Input shapes - lm: {lm.shape}, lm3D: {lm3D.shape}") if hasattr(lm, 'shape') else None
w0, h0 = img.size
print(f"[DEBUG] Original image size: {w0}x{h0}")
# Extract 5 facial landmarks
if lm.shape[0] != 5:
lm5p = extract_5p(lm)
else:
lm5p = lm
print(f"[DEBUG] Landmark points shape: {lm5p.shape}")
# Calculate translation and scale factors
t, s = POS(lm5p.transpose(), lm3D.transpose())
s = rescale_factor / s
# Ensure t is a flat numpy array with exactly 2 elements
t = np.array(t).flatten()
if len(t) != 2:
raise ValueError(f"Translation vector t should have 2 elements, got {len(t)}: {t}")
print(f"[DEBUG] Calculated values - t: {t}, s: {s}")
# Process the image
img_new, lm_new, mask_new = resize_n_crop_img(
img, lm, t, s, target_size=target_size, mask=mask
)
# Create transformation parameters with type checking
trans_params = np.array([
float(w0), # Convert to float explicitly
float(h0), # Convert to float explicitly
float(s), # Convert to float explicitly
float(t[0]), # First translation component
float(t[1]) # Second translation component
], dtype=np.float32)
print(f"[DEBUG] Transformation params: {trans_params}")
return trans_params, img_new, lm_new, mask_new
except Exception as e:
print(f"\n[ERROR] in align_img(): {str(e)}")
print("[DEBUG] Problem occurred with:")
print(f"- img size: {img.size if img else 'None'}")
print(f"- lm shape: {lm.shape if hasattr(lm, 'shape') else 'Not an array'}")
print(f"- lm3D shape: {lm3D.shape if hasattr(lm3D, 'shape') else 'Not an array'}")
print(f"- t: {t if 't' in locals() else 'Not calculated'}")
print(f"- s: {s if 's' in locals() else 'Not calculated'}")
raise
|