-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
executable file
·99 lines (81 loc) · 3.7 KB
/
main.py
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
#!/usr/bin/python3
import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import fitting
import landmarks
FRAMES_PATH = sys.argv[1]
DEF_SHAPE = 'defShape'
OUTPUT_DIR_3D_FINAL = 'final_deformation_heatmaps'
OUTPUT_DIR_3D_INTERMEDIATE = 'intermediate_deformation_heatmaps'
OUTPUT_DIR_PLOT = 'temporal_deformation_plots'
N = 4 # number of frames analysed in each sequence
MEDIAN_ITER = 1 # number of iteration to get median
VMAX = 4 # saturation value for uniform 3d model scaling
LR = True # set to True if using the LR dataset (due to additional initial frames issue)
for directory in {OUTPUT_DIR_3D_FINAL, OUTPUT_DIR_3D_INTERMEDIATE, landmarks.OUTPUT_DIR_2D, OUTPUT_DIR_PLOT}:
if not os.access(directory, os.F_OK):
os.mkdir(directory)
fig_all, ax_all = plt.subplots()
ax_all.set_xlabel('frame fraction')
ax_all.set_ylabel('deformation index')
for i in range(1, len(os.listdir(FRAMES_PATH)) + 1):
path = f'{FRAMES_PATH}/{i:04}'
start = int(
os.path.splitext(min(os.listdir(path), key=lambda s: int(os.path.splitext(s)[0])))[0])
if LR:
start += i - 1
def_shapes = []
def_shapes0 = []
for k in range(start, start + MEDIAN_ITER):
img_path = f'{path}/{k}.jpg'
preds0 = landmarks.lm_dir(img_path, True)
def_shape00 = fitting.fit_3dmm(preds0)[DEF_SHAPE]
def_shape0 = def_shape00.copy()
def_shape00 = fitting.fit_3dmm(preds0, def_shape00)[DEF_SHAPE]
def_shapes.append(def_shape00)
def_shapes0.append(def_shape0)
def_shapes_a = np.array(def_shapes)
def_shape00 = np.median(def_shapes_a, axis=0)
def_shapes0_a = np.array(def_shapes0)
def_shape0 = np.median(def_shapes0_a, axis=0)
y = [0]
fig_intermediate = plt.figure(figsize=plt.figaspect(N))
for j in range(1, N + 1):
def_shapes = []
end = int(os.path.splitext(max(os.listdir(path), key=lambda s: int(os.path.splitext(s)[0])))[0])
n = start + int((end - start) * j / N)
for k in range(n + 1 - MEDIAN_ITER, n + 1):
img_path = f'{path}/{k}.jpg'
preds1 = landmarks.lm_dir(img_path, True)
def_shape01 = def_shape0.copy()
def_shape01 = fitting.fit_3dmm(preds1, def_shape01)[DEF_SHAPE]
def_shapes.append(def_shape01)
def_shapes_a = np.array(def_shapes)
def_shape_median = np.median(def_shapes_a, axis=0)
def_indices = landmarks.models_deformation(def_shape00, def_shape_median)