dockerfile and inference scripts

.dockerignore

@@ -0,0 +1,2 @@
+**
+!ckpts/

dockerfile

@@ -0,0 +1,23 @@
+FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel
+RUN mkdir /workspace/MakeItTalk
+WORKDIR /workspace/MakeItTalk
+
+
+# uncomment this if you have trouble downloading packages
+# RUN rm /etc/apt/sources.list.d/cuda.list && \
+#     rm /etc/apt/sources.list.d/nvidia-ml.list && \
+#     sed -i "s@http://.*archive.ubuntu.com@https://mirrors.tuna.tsinghua.edu.cn@g" /etc/apt/sources.list && \
+#     sed -i "s@http://.*security.ubuntu.com@https://mirrors.tuna.tsinghua.edu.cn@g" /etc/apt/sources.list && \
+#     pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+
+RUN apt-get -y update --fix-missing && \
+    apt-get install -y libglib2.0-0 libxrender1 ffmpeg libsm6 && \
+    pip install ffmpeg-python opencv-python face_alignment scikit-learn pydub pynormalize soundfile librosa pysptk pyworld resemblyzer tensorboardX matplotlib
+
+# Downloading: "https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth" to ckpts/s3fd-619a316812.pth
+# Downloading: "https://www.adrianbulat.com/downloads/python-fan/3DFAN4-4a694010b9.zip" to ckpts/3DFAN4-4a694010b9.zip
+# Downloading: "https://www.adrianbulat.com/downloads/python-fan/depth-6c4283c0e0.zip" to ckpts/checkpoints/depth-6c4283c0e0.zip
+RUN mkdir -p  /root/.cache/torch/hub/checkpoints
+COPY ckpts/s3fd-619a316812.pth /root/.cache/torch/hub/checkpoints/s3fd-619a316812.pth
+COPY ckpts/3DFAN4-4a694010b9.zip /root/.cache/torch/hub/checkpoints/3DFAN4-4a694010b9.zip
+COPY ckpts/depth-6c4283c0e0.zip /root/.cache/torch/hub/checkpoints/depth-6c4283c0e0.zip

inference.sh

@@ -0,0 +1,19 @@
+# example of inference
+base_dir='/mnt/disk_1/xiaoke/codes/talking_head/Wav2Lip/temp/data'
+output_dir='/mnt/disk_1/xiaoke/codes/talking_head/MakeItTalk/examples'
+save_dir='/mnt/disk_1/xiaoke/codes/talking_head/MakeItTalk/examples/outputs'
+mkdir -p $save_dir
+for i in `ls $base_dir`; do
+    face_path="$base_dir/$i/val.png"
+    audio_path="$base_dir/$i/aud.val.wav"
+
+    output_face_path="$output_dir/$i.jpg"
+    output_audio_path="$output_dir/$i.wav"
+    convert $face_path -resize 256x256 $output_face_path
+    cp $audio_path $output_audio_path
+
+    docker run  --gpus all -e CUDA_VISIBLE_DEVICES=5 -it --rm --name makeittalk -v $(realpath .):/workspace/MakeItTalk --shm-size 8gb  makeittalk  python main_end2end.py --jpg "$(basename $output_face_path)"
+    mv $output_dir/*.mp4 $save_dir
+    rm -f $output_face_path $output_audio_path
+done
+

main_end2end.py

@@ -96,7 +96,7 @@
 au_data = []
 au_emb = []
 ains = glob.glob1('examples', '*.wav')
-ains = [item for item in ains if item is not 'tmp.wav']
+ains = [item for item in ains if item != 'tmp.wav']
 ains.sort()
 for ain in ains:
     os.system('ffmpeg -y -loglevel error -i examples/{} -ar 16000 examples/tmp.wav'.format(ain))

src/approaches/train_image_translation.py

@@ -383,7 +383,10 @@ def single_test(self, jpg=None, fls=None, filename=None, prefix='', grey_only=Fa
             fls[:, 0::3] += 130
             fls[:, 1::3] += 80
 
-        writer = cv2.VideoWriter('out.mp4', cv2.VideoWriter_fourcc(*'mjpg'), 62.5, (256 * 3, 256))
+        writer = cv2.VideoWriter('out.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 62.5, (256 * 3, 256))
+        video_writer = cv2.VideoWriter(
+            'examples/{}_{}.output.mp4'.format(prefix, filename[:-4]), cv2.VideoWriter_fourcc(*'mp4v'), 62.5, (256, 256)
+        )
 
         for i, frame in enumerate(fls):
 
@@ -420,16 +423,19 @@ def single_test(self, jpg=None, fls=None, filename=None, prefix='', grey_only=Fa
             for i in range(g_out.shape[0]):
                 frame = np.concatenate((ref_in[i], g_out[i], fls_in[i]), axis=1) * 255.0
                 writer.write(frame.astype(np.uint8))
+                video_writer.write((g_out[i] * 255.0).astype(np.uint8))
 
         writer.release()
+        video_writer.release()
         print('Time - only video:', time.time() - st)
 
         if(filename is None):
             filename = 'v'
         os.system('ffmpeg -loglevel error -y -i out.mp4 -i {} -pix_fmt yuv420p -strict -2 examples/{}_{}.mp4'.format(
             'examples/'+filename[9:-16]+'.wav',
             prefix, filename[:-4]))
-        # os.system('rm out.mp4')
+        os.system('rm out.mp4')
+        os.system("ffmpeg -loglevel error -y -i examples/{}_{}.output.mp4 -i {} -pix_fmt yuv420p -strict -2 examples/{}_{}.output.audio.mp4".format(prefix, filename[:-4], 'examples/'+filename[9:-16]+'.wav', prefix, filename[:-4]))
 
         print('Time - ffmpeg add audio:', time.time() - st)