def buildTrainData(baseDir):
model = load_fan()
dirs = [baseDir + '/' + d for d in os.listdir(baseDir)]
train_str = ''
for d in dirs:
mainPic = d + '/main.jpg'
pics = getPicList(d)
for p in pics:
img = cv2.imread(p, cv2.IMREAD_COLOR)
points = get_68_points(img, model)
if points is None:
continue
px = points[:, 0]
py = points[:, 1]
x_str = ''
y_str = ''
for i in range(points.shape[0]):
x_str += str(int(px[i])) + ' '
y_str += str(int(py[i])) + ' '
x_str = x_str.strip()
y_str = y_str.strip()
p_str = p + ' ' + mainPic + ' ' + x_str + ' ' + y_str
train_str += p_str + '\n'
print(p)
open('train_data.txt', 'w', 'utf-8').write(train_str.strip())
def generateAB68(pathA, pathB):
Afiles = getFileList(pathA)
Bfiles = getFileList(pathB)
A_str = ''
B_str = ''
model = load_fan()
for a in Afiles:
m_str = innerGen68(a, model)
if m_str is not None:
A_str += a + ' ' + m_str + '\n'
print(a, m_str)
for b in Bfiles:
m_str = innerGen68(b, model)
if m_str is not None:
B_str += b + ' ' + m_str + '\n'
print(b, m_str)
open('A_68.txt', 'w', 'utf-8').write(A_str.strip())
open('B_68.txt', 'w', 'utf-8').write(B_str.strip())
def buildSingleData(baseDir):
model = load_fan()
train_str = ''
paths = [baseDir + '/' + p for p in os.listdir(baseDir)]
for p in paths:
img = cv2.imread(p, cv2.IMREAD_COLOR)
points = get_68_points(img, model)
if points is None:
continue
px = points[:, 0]
py = points[:, 1]
x_str = ''
y_str = ''
for i in range(points.shape[0]):
x_str += str(int(px[i])) + ' '
y_str += str(int(py[i])) + ' '
x_str = x_str.strip()
y_str = y_str.strip()
p_str = p + ' ' + x_str + ' ' + y_str
train_str += p_str + '\n'
print(p)
open('train_data.txt', 'w', 'utf-8').write(train_str.strip())
def processSingleMain(path):
model = load_fan()
mp = path + '/main.jpg'
img = getMainFace(mp, model)
cv2.imwrite(mp, img)
print('write', mp)
import os, sys
sys.path.append('./face-alignment')
from MyFanPred import get_68_points as get_68_fan_points
from MyFanPred import load_fan
from mtcnn import DetectFace
from skimage import transform as trans
from color_adjust import Color, MatchHist
from fan_util import drawLandmarks
from SAE_clip_resnext_dln import AutoEncoder
sys.path.append('./resoaugment')
from next_resoaug import SAE_RESNEXT_ENCODER as RESOAUGMODEL
sys.path.append('./CosFace_pytorch')
from cosface_pred import get_img_feature, get_distance
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
fan = load_fan()
color_adj = Color()
detectFace = DetectFace()
AEModel = AutoEncoder()
AEModel.load_state_dict(
torch.load('models/angelababy_libingbing.pth', map_location=device))
AEModel.eval()
resoModel = RESOAUGMODEL()
resoModel.load_state_dict(
torch.load('resoaugment/best_generator.pth', map_location=device))
resoModel.eval()
landmarks_2D = np.array(
[
[0.000213256, 0.106454], #17
[0.0752622, 0.038915], #18