| """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) |
|
|
|
|
| |
| 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 |
| |
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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: |
| |
| 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}") |
|
|
| |
| if lm.shape[0] != 5: |
| lm5p = extract_5p(lm) |
| else: |
| lm5p = lm |
| |
| print(f"[DEBUG] Landmark points shape: {lm5p.shape}") |
|
|
| |
| t, s = POS(lm5p.transpose(), lm3D.transpose()) |
| s = rescale_factor / s |
| |
| |
| 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}") |
|
|
| |
| img_new, lm_new, mask_new = resize_n_crop_img( |
| img, lm, t, s, target_size=target_size, mask=mask |
| ) |
|
|
| |
| trans_params = np.array([ |
| float(w0), |
| float(h0), |
| float(s), |
| float(t[0]), |
| float(t[1]) |
| ], 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 |
|
|
