File size: 3,191 Bytes
03da825 |
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 |
import torch
import torchvision.transforms as transforms
import cv2
import face_alignment
import numpy as np
from PIL import Image
from data_preprocessing.lm_based_image_align import LandmarkBasedImageAlign
from data_preprocessing.pose_estimation import PoseEstimator
from data_preprocessing.crop_image import ImageCropper
from data_preprocessing.foreground_extractor import ForegroundExtractor
class DataPreprocessor:
def __init__(self, device, crop_smooth_alpha=0.9):
self.device = device
self.crop_smooth_alpha = crop_smooth_alpha
self.face_alignment = LandmarkBasedImageAlign(output_size=1024, transform_size=1024)
self.pose_estimator = PoseEstimator(device)
self.cropper = ImageCropper()
self.foreground_extractor = ForegroundExtractor(device)
self.fa = face_alignment.FaceAlignment(
face_alignment.LandmarksType.TWO_D, flip_input=False, device=device
)
self.transform = transforms.ToTensor()
self.lm_woAlign_ema = None
self.aligned_lm_ema = None
def __call__(self, img, keep_bg=False):
lm = self.fa.get_landmarks(np.array(img))
if lm is None:
detected_face = [0, 0, img.size[0], img.size[1]]
lm = self.fa.get_landmarks(img, detected_faces=[detected_face])[0]
else:
lm = lm[0]
if self.lm_woAlign_ema is None:
self.lm_woAlign_ema = np.array(lm)
self.lm_woAlign_ema = (self.lm_woAlign_ema * 0.9 + np.array(lm) * 0.1)
img_woAlign, lm_woAlign = self.cropper(img, self.lm_woAlign_ema)
lm_woAlign = np.array(lm_woAlign)
img_woAlign = np.array(img_woAlign)
img, aligned_lm = self.face_alignment(img, lm)
intrinsics, pose = self.pose_estimator.predict_pose(img, aligned_lm)
img, aligned_lm = self.cropper(img, aligned_lm)
img = np.array(img)
if not keep_bg:
matte = self.foreground_extractor(img)
img = (img * matte).astype(np.uint8)
if self.aligned_lm_ema is None:
self.aligned_lm_ema = aligned_lm
self.aligned_lm_ema = 0.8 * self.aligned_lm_ema + aligned_lm * 0.2
crop_params = cv2.estimateAffinePartial2D(self.aligned_lm_ema, lm_woAlign)[0]
return img, img_woAlign, intrinsics, pose, crop_params
def from_path(self, image_path, device, keep_bg=False):
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img)
img, _, intrinsics, pose, crop_params = self(img, keep_bg=keep_bg)
img = np.transpose(img, (2, 0, 1))[None, :, :, :] / 255.
img = (img * 2 - 1)
img = torch.from_numpy(img).float()
pose = torch.from_numpy(pose).unsqueeze(0).float()
intrinsics = torch.from_numpy(intrinsics).unsqueeze(0).float()
return {
'image': img.to(device),
'cam2world': pose.to(device),
'intrinsics': intrinsics.to(device),
'crop_params': crop_params
}
@staticmethod
def realign(img, T):
out = cv2.warpAffine(img, T, (img.shape[1], img.shape[0]))
return out
|