{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path \n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "import sys\n",
    "sys.path.append(str(Path(os.path.abspath('')).parent))\n",
    "\n",
    "import torch\n",
    "import torch.distributions as D\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "import wandb\n",
    "from tqdm import tqdm\n",
    "api = wandb.Api()\n",
    "\n",
    "agent_path = Path(os.path.abspath('')).parent / 'models' / 'genrl_stickman_500k_2.pt'\n",
    "print(\"Model path\", agent_path)\n",
    "\n",
    "agent = torch.load(agent_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tools.genrl_utils import ViCLIPGlobalInstance, DOMAIN2PREDICATES\n",
    "model_name = getattr(agent.cfg, 'viclip_model', 'viclip')\n",
    "# Get ViCLIP\n",
    "if 'viclip_global_instance' not in locals() or model_name != viclip_global_instance._model:\n",
    "    viclip_global_instance = ViCLIPGlobalInstance(model_name)\n",
    "    if not viclip_global_instance._instantiated:\n",
    "        print(\"Instantiating\")\n",
    "        viclip_global_instance.instantiate()\n",
    "    clip = viclip_global_instance.viclip\n",
    "    tokenizer = viclip_global_instance.viclip_tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "\n",
    "def get_vid_feat(frames, clip):\n",
    "    return clip.get_vid_features(frames,)\n",
    "\n",
    "def _frame_from_video(video):\n",
    "    while video.isOpened():\n",
    "        success, frame = video.read()\n",
    "        if success:\n",
    "            yield frame\n",
    "        else:\n",
    "            break\n",
    "\n",
    "v_mean = np.array([0.485, 0.456, 0.406]).reshape(1,1,3)\n",
    "v_std = np.array([0.229, 0.224, 0.225]).reshape(1,1,3)\n",
    "def normalize(data):\n",
    "    return (data/255.0-v_mean)/v_std\n",
    "\n",
    "def denormalize(data):\n",
    "    return (((data * v_std) + v_mean) * 255) \n",
    "\n",
    "def frames2tensor(vid_list, fnum=8, target_size=(224, 224), device=torch.device('cuda')):\n",
    "    vid_list = [*vid_list[0]]\n",
    "    assert(len(vid_list) >= fnum)\n",
    "    vid_list = [cv2.resize(x, target_size) for x in vid_list]\n",
    "    vid_tube = [np.expand_dims(normalize(x), axis=(0, 1)) for x in vid_list]\n",
    "    vid_tube = np.concatenate(vid_tube, axis=1)\n",
    "    vid_tube = np.transpose(vid_tube, (0, 1, 4, 2, 3))\n",
    "    vid_tube = torch.from_numpy(vid_tube).to(device, non_blocking=True).float()\n",
    "    return vid_tube\n",
    "\n",
    "\n",
    "def get_video_feat(frames, device=torch.device('cuda'), flip=False):\n",
    "    # Image\n",
    "    if frames.shape[1] == 1:\n",
    "        frames = frames.transpose(1,0,2,3,4).repeat(8, axis=0).transpose(1,0,2,3,4)\n",
    "\n",
    "    # Short video\n",
    "    if frames.shape[1] == 4:\n",
    "        frames = frames.transpose(1,0,2,3,4).repeat(2, axis=0).transpose(1,0,2,3,4)\n",
    "\n",
    "    k = max(frames.shape[1] // 128, 1)\n",
    "    frames = frames[:, ::k]\n",
    "    \n",
    "    # Horizontally flip\n",
    "    if flip:\n",
    "        frames = np.flip(frames, axis=-2)\n",
    "\n",
    "    print(frames.shape,)\n",
    "    chosen_frames = frames[:, :8]\n",
    "    chosen_frames = frames2tensor(chosen_frames, device=device)\n",
    "    vid_feat = get_vid_feat(chosen_frames, clip,)\n",
    "    return vid_feat, chosen_frames\n",
    "\n",
    "VIDEO_PATH = Path(os.path.abspath('')).parent / 'assets' / 'video_samples'\n",
    "video_name = 'headstand.mp4'\n",
    "\n",
    "video_file_path = str(VIDEO_PATH / video_name)\n",
    "print(video_file_path)\n",
    "video = cv2.VideoCapture(video_file_path)\n",
    "frames = np.expand_dims(np.stack([ cv2.cvtColor(x, cv2.COLOR_BGR2RGB) for x in _frame_from_video(video)], axis=0), axis=0)\n",
    "print('Video length:', frames.shape[1])\n",
    "with torch.no_grad():\n",
    "    vid_feat, frames_feat = get_video_feat(frames, flip=False)\n",
    "print(vid_feat.shape)\n",
    "plt.imshow(frames[0,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "video_embed = vid_feat\n",
    "DENOISE = True\n",
    "\n",
    "T = video_embed.shape[0]\n",
    "\n",
    "from torchvision.transforms import transforms as vision_trans\n",
    "trasnf = vision_trans.Resize(size=(64, 64), interpolation=vision_trans.InterpolationMode.NEAREST)\n",
    "\n",
    "wm = world_model = agent.wm\n",
    "connector = agent.wm.connector\n",
    "decoder = world_model.heads['decoder']\n",
    "n_frames = connector.n_frames\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    # Get actions\n",
    "    video_embed = video_embed.unsqueeze(1).repeat(1,n_frames, 1).reshape(1, n_frames * T, -1)\n",
    "    action = wm.connector.get_action(video_embed)\n",
    "\n",
    "    # Imagine\n",
    "    prior = wm.connector.video_imagine(video_embed, None, sample=False, reset_every_n_frames=False, denoise=DENOISE)\n",
    "    prior_recon = decoder(wm.decoder_input_fn(prior))['observation'].mean + 0.5\n",
    "\n",
    "    # Plotting video\n",
    "    ims = []\n",
    "    fig, axes = plt.subplots(1, 1, figsize=(4, 8), frameon=False)\n",
    "    fig.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)\n",
    "    fig.set_size_inches(4,2)\n",
    "\n",
    "    for t in range(prior_recon.shape[1]):\n",
    "        toadd = []\n",
    "        for b in range(prior_recon.shape[0]):\n",
    "            ax = axes\n",
    "            ax.set_axis_off()\n",
    "            img = cv2.resize((np.clip(prior_recon[b, t].cpu().permute(1,2,0), 0, 1).numpy() *255).astype(np.uint8), (224,224))\n",
    "            orig_img = denormalize(frames_feat[b, t].cpu().permute(1,2,0) ).numpy().astype(np.uint8)\n",
    "            frame =  ax.imshow(np.concatenate([orig_img, img], axis=1))   \n",
    "            toadd.append(frame) # add both the image and the text to the list of artists \n",
    "        ims.append(toadd)\n",
    "\n",
    "    anim = animation.ArtistAnimation(fig, ims, interval=700, blit=True, repeat_delay=700, )\n",
    "\n",
    "    # Save GIFs\n",
    "    writer = animation.PillowWriter(fps=15, metadata=dict(artist='Me'), bitrate=1800,)\n",
    "    domain = agent.cfg.task.split('_')[0]\n",
    "    os.makedirs(f'videos/{domain}/video2video', exist_ok=True)\n",
    "    file_path = f'videos/{domain}/video2video/{video_name[:-4].replace(\" \",\"_\")}.gif'\n",
    "    anim.save(file_path, writer=writer, )\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('base')",
   "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.10.14"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "3d597f4c481aa0f25dceb95d2a0067e73c0966dcbd003d741d821a7208527ecf"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}