first-order-model-demo (2024).ipynb

{"cells":[{"cell_type":"code","execution_count":null,"metadata":{"id":"ES3TyCd6xFM6"},"outputs":[],"source":["from IPython.display import HTML\n","from demo import load_checkpoints\n","from demo import make_animation\n","from skimage import img_as_ubyte\n","from skimage.transform import resize\n","import imageio\n","import matplotlib.animation as animation\n","import matplotlib.pyplot as plt\n","import numpy as np\n","import warnings"]},{"cell_type":"code","execution_count":10,"metadata":{"id":"RyvHtAErxNQK"},"outputs":[],"source":["def display(source, driving, generated=None):\n","    fig = plt.figure(figsize=(8 + 4 * (generated is not None), 6))\n","    ims = []\n","    for i in range(len(driving)):\n","        cols = [source]\n","        cols.append(driving[i])\n","        if generated is not None:\n","            cols.append(generated[i])\n","        im = plt.imshow(np.concatenate(cols, axis=1), animated=True)\n","        plt.axis(\"off\")\n","        ims.append([im])\n","    ani = animation.ArtistAnimation(fig, ims, interval=50, repeat_delay=1000)\n","    plt.close()\n","    return ani"]},{"cell_type":"code","execution_count":12,"metadata":{"id":"Z9553Y1DxMV3"},"outputs":[],"source":["warnings.filterwarnings(\"ignore\")\n","source_image = imageio.imread(\"demo/images/01.png\")\n","reader = imageio.get_reader(\"demo/videos/0.mp4\")\n","source_image = resize(source_image, (256, 256))[..., :3]\n","fps = reader.get_meta_data()[\"fps\"]\n","driving_video = []\n","try:\n","    for im in reader:\n","        driving_video.append(im)\n","except RuntimeError:\n","    pass\n","reader.close()\n","driving_video = [resize(frame, (256, 256))[..., :3] for frame in driving_video]"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Oxi6-riLOgnm"},"outputs":[],"source":["\"\"\"\n","HTML(display(source_image, driving_video).to_html5_video())\n","generator, kp_detector = load_checkpoints(\n","    config_path=\"config/vox-256.yaml\", checkpoint_path=\"vox-cpk.pth.tar\"\n",")\n","\"\"\"\n","generator, kp_detector = load_checkpoints(\n","    config_path=\"config/vox-adv-256.yaml\", checkpoint_path=\"vox-adv-cpk.pth.tar\"\n",")\n","predictions = make_animation(\n","    source_image, driving_video, generator, kp_detector, relative=True\n",")\n","imageio.mimsave(\n","    \"generated.mp4\", [img_as_ubyte(frame) for frame in predictions], fps=fps\n",")\n","HTML(display(source_image, driving_video, predictions).to_html5_video())"]}],"metadata":{"accelerator":"GPU","kernelspec":{"display_name":"fomd","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8.20"}},"nbformat":4,"nbformat_minor":0}