Spaces:
Running
on
A10G
Running
on
A10G
File size: 6,459 Bytes
a22eb82 a86a2b8 a22eb82 8cc4c88 a22eb82 a86a2b8 a22eb82 0ce42bd a22eb82 416263d a22eb82 0ce42bd a86a2b8 a22eb82 416263d 0ce42bd a22eb82 0ce42bd a86a2b8 a22eb82 a86a2b8 a22eb82 0ce42bd a22eb82 0ce42bd |
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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
import numpy as np
import cv2, os, sys, torch
from tqdm import tqdm
from PIL import Image
# 3dmm extraction
from src.face3d.util.preprocess import align_img
from src.face3d.util.load_mats import load_lm3d
from src.face3d.models import networks
from src.face3d.extract_kp_videos import KeypointExtractor
from scipy.io import loadmat, savemat
from src.utils.croper import Croper
import warnings
warnings.filterwarnings("ignore")
def split_coeff(coeffs):
"""
Return:
coeffs_dict -- a dict of torch.tensors
Parameters:
coeffs -- torch.tensor, size (B, 256)
"""
id_coeffs = coeffs[:, :80]
exp_coeffs = coeffs[:, 80: 144]
tex_coeffs = coeffs[:, 144: 224]
angles = coeffs[:, 224: 227]
gammas = coeffs[:, 227: 254]
translations = coeffs[:, 254:]
return {
'id': id_coeffs,
'exp': exp_coeffs,
'tex': tex_coeffs,
'angle': angles,
'gamma': gammas,
'trans': translations
}
class CropAndExtract():
def __init__(self, path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, device):
self.croper = Croper(path_of_lm_croper)
self.kp_extractor = KeypointExtractor(device)
self.net_recon = networks.define_net_recon(net_recon='resnet50', use_last_fc=False, init_path='').to(device)
checkpoint = torch.load(path_of_net_recon_model, map_location=torch.device(device))
self.net_recon.load_state_dict(checkpoint['net_recon'])
self.net_recon.eval()
self.lm3d_std = load_lm3d(dir_of_BFM_fitting)
self.device = device
def generate(self, input_path, save_dir, crop_or_resize='crop'):
pic_size = 256
pic_name = os.path.splitext(os.path.split(input_path)[-1])[0]
landmarks_path = os.path.join(save_dir, pic_name+'_landmarks.txt')
coeff_path = os.path.join(save_dir, pic_name+'.mat')
png_path = os.path.join(save_dir, pic_name+'.png')
#load input
if not os.path.isfile(input_path):
raise ValueError('input_path must be a valid path to video/image file')
elif input_path.split('.')[-1] in ['jpg', 'png', 'jpeg']:
# loader for first frame
full_frames = [cv2.imread(input_path)]
fps = 25
else:
# loader for videos
video_stream = cv2.VideoCapture(input_path)
fps = video_stream.get(cv2.CAP_PROP_FPS)
full_frames = []
while 1:
still_reading, frame = video_stream.read()
if not still_reading:
video_stream.release()
break
full_frames.append(frame)
x_full_frames= [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in full_frames]
#### crop images as the
if crop_or_resize.lower() == 'crop': # default crop
x_full_frames, crop, quad = self.croper.crop(x_full_frames, xsize=pic_size)
clx, cly, crx, cry = crop
lx, ly, rx, ry = quad
lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)
oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx
crop_info = ((ox2 - ox1, oy2 - oy1), crop, quad)
elif crop_or_resize.lower() == 'full':
x_full_frames, crop, quad = self.croper.crop(x_full_frames, still=True, xsize=pic_size)
clx, cly, crx, cry = crop
lx, ly, rx, ry = quad
lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)
oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx
crop_info = ((ox2 - ox1, oy2 - oy1), crop, quad)
else: # resize mode
oy1, oy2, ox1, ox2 = 0, x_full_frames[0].shape[0], 0, x_full_frames[0].shape[1]
crop_info = ((ox2 - ox1, oy2 - oy1), None, None)
frames_pil = [Image.fromarray(cv2.resize(frame,(pic_size, pic_size))) for frame in x_full_frames]
if len(frames_pil) == 0:
print('No face is detected in the input file')
return None, None
# save crop info
for frame in frames_pil:
cv2.imwrite(png_path, cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR))
# 2. get the landmark according to the detected face.
if not os.path.isfile(landmarks_path):
lm = self.kp_extractor.extract_keypoint(frames_pil, landmarks_path)
else:
print(' Using saved landmarks.')
lm = np.loadtxt(landmarks_path).astype(np.float32)
lm = lm.reshape([len(x_full_frames), -1, 2])
if not os.path.isfile(coeff_path):
# load 3dmm paramter generator from Deep3DFaceRecon_pytorch
video_coeffs, full_coeffs = [], []
for idx in tqdm(range(len(frames_pil)), desc='3DMM Extraction In Video:'):
frame = frames_pil[idx]
W,H = frame.size
lm1 = lm[idx].reshape([-1, 2])
if np.mean(lm1) == -1:
lm1 = (self.lm3d_std[:, :2]+1)/2.
lm1 = np.concatenate(
[lm1[:, :1]*W, lm1[:, 1:2]*H], 1
)
else:
lm1[:, -1] = H - 1 - lm1[:, -1]
trans_params, im1, lm1, _ = align_img(frame, lm1, self.lm3d_std)
trans_params = np.array([float(item) for item in np.hsplit(trans_params, 5)]).astype(np.float32)
im_t = torch.tensor(np.array(im1)/255., dtype=torch.float32).permute(2, 0, 1).to(self.device).unsqueeze(0)
with torch.no_grad():
full_coeff = self.net_recon(im_t)
coeffs = split_coeff(full_coeff)
pred_coeff = {key:coeffs[key].cpu().numpy() for key in coeffs}
pred_coeff = np.concatenate([
pred_coeff['exp'],
pred_coeff['angle'],
pred_coeff['trans'],
trans_params[2:][None],
], 1)
video_coeffs.append(pred_coeff)
full_coeffs.append(full_coeff.cpu().numpy())
semantic_npy = np.array(video_coeffs)[:,0]
savemat(coeff_path, {'coeff_3dmm': semantic_npy, 'full_3dmm': np.array(full_coeffs)[0]})
return coeff_path, png_path, crop_info
|