Visualization-Copy2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paper visualizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --user neural_renderer_pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import imageio\n",
    "import trimesh\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.cm as cm\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import neural_renderer as nr\n",
    "from scipy.spatial import cKDTree as KDTree\n",
    "from datasets import make_data_instance_from_stl\n",
    "\n",
    "from models import *\n",
    "\n",
    "import pdb\n",
    "\n",
    "def get_rotate_matrix(rotation_angle1):\n",
    "    cosval = np.cos(rotation_angle1)\n",
    "    sinval = np.sin(rotation_angle1)\n",
    "\n",
    "    rotation_matrix_x = np.array([[1, 0, 0, 0],\n",
    "                                  [0, cosval, -sinval, 0],\n",
    "                                  [0, sinval, cosval, 0],\n",
    "                                  [0, 0, 0, 1]])\n",
    "    rotation_matrix_y = np.array([[cosval, 0, sinval, 0],\n",
    "                                  [0, 1, 0, 0],\n",
    "                                  [-sinval, 0, cosval, 0],\n",
    "                                  [0, 0, 0, 1]])\n",
    "    rotation_matrix_z = np.array([[cosval, -sinval, 0, 0],\n",
    "                                  [sinval, cosval, 0, 0],\n",
    "                                  [0, 0, 1, 0],\n",
    "                                  [0, 0, 0, 1]])\n",
    "    scale_y_neg = np.array([\n",
    "        [1, 0, 0, 0],\n",
    "        [0, -1, 0, 0],\n",
    "        [0, 0, 1, 0],\n",
    "        [0, 0, 0, 1]\n",
    "    ])\n",
    "\n",
    "    neg = np.array([\n",
    "        [-1, 0, 0, 0],\n",
    "        [0, -1, 0, 0],\n",
    "        [0, 0, -1, 0],\n",
    "        [0, 0, 0, 1]\n",
    "    ])\n",
    "    # y,z swap = x rotate -90, scale y -1\n",
    "    # new_pts0[:, 1] = new_pts[:, 2]\n",
    "    # new_pts0[:, 2] = new_pts[:, 1]\n",
    "    #\n",
    "    # x y swap + negative = z rotate -90, scale y -1\n",
    "    # new_pts0[:, 0] = - new_pts0[:, 1] = - new_pts[:, 2]\n",
    "    # new_pts0[:, 1] = - new_pts[:, 0]\n",
    "\n",
    "    # return np.linalg.multi_dot([rotation_matrix_z, rotation_matrix_y, rotation_matrix_y, scale_y_neg, rotation_matrix_z, scale_y_neg, rotation_matrix_x])\n",
    "    return np.linalg.multi_dot([neg, rotation_matrix_z, rotation_matrix_z, scale_y_neg, rotation_matrix_x])\n",
    "\n",
    "def get_projection_matricies(az, el, distance_ratio, roll = 0, focal_length=35, img_w=137, img_h=137):\n",
    "    \"\"\"\n",
    "    Calculate 4x3 3D to 2D projection matrix given viewpoint parameters.\n",
    "    Code from \"https://github.com/Xharlie/DISN\"\n",
    "    \"\"\"\n",
    "\n",
    "    F_MM = focal_length  # Focal length\n",
    "    SENSOR_SIZE_MM = 32.\n",
    "    PIXEL_ASPECT_RATIO = 1.  # pixel_aspect_x / pixel_aspect_y\n",
    "    RESOLUTION_PCT = 100.\n",
    "    SKEW = 0.\n",
    "    CAM_MAX_DIST = 1.75\n",
    "    CAM_ROT = np.asarray([[1.910685676922942e-15, 4.371138828673793e-08, 1.0],\n",
    "                      [1.0, -4.371138828673793e-08, -0.0],\n",
    "                      [4.371138828673793e-08, 1.0, -4.371138828673793e-08]])\n",
    "\n",
    "    # Calculate intrinsic matrix.\n",
    "    scale = RESOLUTION_PCT / 100\n",
    "    # print('scale', scale)\n",
    "    f_u = F_MM * img_w * scale / SENSOR_SIZE_MM\n",
    "    f_v = F_MM * img_h * scale * PIXEL_ASPECT_RATIO / SENSOR_SIZE_MM\n",
    "    # print('f_u', f_u, 'f_v', f_v)\n",
    "    u_0 = img_w * scale / 2\n",
    "    v_0 = img_h * scale / 2\n",
    "    K = np.matrix(((f_u, SKEW, u_0), (0, f_v, v_0), (0, 0, 1)))\n",
    "\n",
    "    # Calculate rotation and translation matrices.\n",
    "    # Step 1: World coordinate to object coordinate.\n",
    "    sa = np.sin(np.radians(-az))\n",
    "    ca = np.cos(np.radians(-az))\n",
    "    se = np.sin(np.radians(-el))\n",
    "    ce = np.cos(np.radians(-el))\n",
    "    R_world2obj = np.transpose(np.matrix(((ca * ce, -sa, ca * se),\n",
    "                                          (sa * ce, ca, sa * se),\n",
    "                                          (-se, 0, ce))))\n",
    "\n",
    "    # Step 2: Object coordinate to camera coordinate.\n",
    "    R_obj2cam = np.transpose(np.matrix(CAM_ROT))\n",
    "    R_world2cam = R_obj2cam * R_world2obj\n",
    "    cam_location = np.transpose(np.matrix((distance_ratio * CAM_MAX_DIST,\n",
    "                                           0,\n",
    "                                           0)))\n",
    "    T_world2cam = -1 * R_obj2cam * cam_location\n",
    "\n",
    "    # Step 3: Fix blender camera's y and z axis direction.\n",
    "    R_camfix = np.matrix(((1, 0, 0), (0, -1, 0), (0, 0, -1)))\n",
    "    R_world2cam = R_camfix * R_world2cam\n",
    "    T_world2cam = R_camfix * T_world2cam\n",
    "\n",
    "    RT = np.hstack((R_world2cam, T_world2cam))\n",
    "    # finally, consider roll\n",
    "    cr = np.cos(np.radians(roll))\n",
    "    sr = np.sin(np.radians(roll))\n",
    "    R_z = np.matrix(((cr, -sr, 0),\n",
    "                  (sr, cr, 0),\n",
    "                  (0, 0, 1)))\n",
    "\n",
    "    rot_mat = get_rotate_matrix(-np.pi / 2)\n",
    "\n",
    "    return K, R_z@RT@rot_mat\n",
    "\n",
    "def load_fld(fld_path):\n",
    "    '''\n",
    "     Takes a path to generated fld file with following colomns: x,y,z,p,k,omega,nut\n",
    "                and converts it into a geometric data instance.\n",
    "    '''\n",
    "\n",
    "    fld = np.genfromtxt(fld_path, delimiter=',', skip_header=1)\n",
    "    np.random.shuffle(fld)\n",
    "    fld[fld > 10e5] = np.nan\n",
    "    fld = fld[~np.isnan(fld).any(axis=1)]\n",
    "    answers = fld[:, 3:]\n",
    "\n",
    "    \"\"\"\n",
    "    mean_values = [-2.06707869e+00, 1.04133005e-01, 2.17513919e+02, 6.04485806e-05]\n",
    "    std_values = [3.71674873e+00, 4.93675056e-02, 1.10871494e+02, 2.63155496e-05]\n",
    "    for f in range(answers.shape[1]):\n",
    "        answers[:, f] = (answers[:, f] - mean_values[f]) / std_values[f]\n",
    "    \"\"\"\n",
    "\n",
    "    stl_path = fld_path.replace('fld', 'stl', 1)[:-9] + '.stl'\n",
    "    mesh = trimesh.load(stl_path)\n",
    "    # reinterpolate features on mesh\n",
    "    fld_tree = KDTree(fld[:, :3])\n",
    "    distances, indeces = fld_tree.query(mesh.vertices, k=1)\n",
    "    interpolations = answers[indeces].squeeze()\n",
    "\n",
    "    return mesh, interpolations\n",
    "\n",
    "def load_predicted(ply_path):\n",
    "    '''\n",
    "     Takes a path to generated fld file with following colomns: x,y,z,p,k,omega,nut\n",
    "                and converts it into a geometric data instance.\n",
    "    '''\n",
    "\n",
    "    answers_path = ply_path.replace('meshes', 'predictions', 1)[:-4] + '.npy'\n",
    "    answers = np.load(answers_path)\n",
    "    mesh = trimesh.load(ply_path)\n",
    "\n",
    "    return mesh, answers\n",
    "\n",
    "\n",
    "def interpolate_on_faces(field, faces):\n",
    "    #TODO: no batch support for now\n",
    "    nv = field.shape[0]\n",
    "    nf = faces.shape[0]\n",
    "    field = field.reshape((nv, 1))\n",
    "    # pytorch only supports long and byte tensors for indexing\n",
    "    face_coordinates = field[faces.long()].squeeze(0)\n",
    "    centroids = 1.0/3 * torch.sum(face_coordinates, 1)\n",
    "    return centroids.squeeze(-1)\n",
    "\n",
    "def visualize(vertices, faces, fields, field_to_visualize = 0, \n",
    "              img_resolution = 1200, azimuth = 210, elevation=10, distance_ratio = 0.8, colormap=cm.jet, \n",
    "              color_blind=False):\n",
    "    \"\"\"\n",
    "    Interface to neural_render to produce nice visualizations. It requires GPU.\n",
    "\n",
    "    Inputs:\n",
    "    vertices in [V,3]\n",
    "    faces in [F,3]\n",
    "    fields in [V,3]\n",
    "    (ideally you can substitute this with a torch_geometric.data.Data object. \n",
    "    I didn't because I don't have it installed)\n",
    "\n",
    "    Output:\n",
    "    Image in [img_resolution, img_resolution, 3]\n",
    "    \"\"\"\n",
    "\n",
    "    # first set up camera\n",
    "    intrinsic, extrinsic =  get_projection_matricies(azimuth, elevation, distance_ratio, img_w=img_resolution, img_h=img_resolution)\n",
    "\n",
    "    K_cuda = torch.tensor(intrinsic[np.newaxis, :, :].copy()).float().cuda().unsqueeze(0)\n",
    "    R_cuda = torch.tensor(extrinsic[np.newaxis, 0:3, 0:3].copy()).float().cuda().unsqueeze(0)\n",
    "    t_cuda = torch.tensor(extrinsic[np.newaxis, np.newaxis, 0:3, 3].copy()).float().cuda().unsqueeze(0)\n",
    "\n",
    "    # initialize renderer\n",
    "    renderer = nr.Renderer(image_size = img_resolution, orig_size = img_resolution, K=K_cuda, R=R_cuda, t=t_cuda, anti_aliasing=True)\n",
    "\n",
    "    # now move vertices, faces to GPU\n",
    "    verts_dr = torch.tensor(vertices.copy(), dtype=torch.float32, requires_grad = False).cuda()\n",
    "    faces_dr = torch.tensor(faces.copy()).cuda()\n",
    "    field_dr = torch.tensor(fields[:, field_to_visualize].copy(),dtype=torch.float32, requires_grad = False).cuda()\n",
    "    # interpolate field on traingle center\n",
    "    field_on_faces = interpolate_on_faces(field_dr, faces_dr)\n",
    "\n",
    "    #TODO: find good values here? Maybe across the dataset to make visualization consistent? or this is good enough? I am not sure...\n",
    "    norm = mpl.colors.Normalize(vmin= -6, vmax=6)\n",
    "    cmap = colormap\n",
    "    m = cm.ScalarMappable(norm=norm, cmap=cmap)\n",
    "    # field_on_faces = torch.clamp((field_on_faces-field_min)/(field_max-field_min),0,1)\n",
    "\n",
    "    textures_dr = torch.ones(faces_dr.shape[0], 1, 1, 1, 3, dtype=torch.float32).cuda()\n",
    "    # feel free to pick your favorite color map here, I used this one for Sanity check, maybe we can  use another one here??\n",
    "    \n",
    "    if not color_blind:\n",
    "        textures_dr[:,0,0,0, :] = torch.tensor(list(map(m.to_rgba,  field_on_faces.cpu().detach())), dtype=torch.float32).cuda()[:, :3]\n",
    "\n",
    "    images_out, depth, alpha = renderer(verts_dr.unsqueeze(0), faces_dr.unsqueeze(0), textures_dr.unsqueeze(0))\n",
    "    images_out = torch.cat([images_out[0], alpha])\n",
    "    image_out_export = 255*images_out.detach().cpu().numpy().transpose((1, 2, 0))\n",
    "    return image_out_export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_data_instance_from_ply(path):\n",
    "    mesh = trimesh.load(path)\n",
    "    edge_attr = [mesh.vertices[a] - mesh.vertices[b] for a, b in mesh.edges]\n",
    "    data = torch_geometric.data.Data(x  = torch.tensor(mesh.vertices, dtype=torch.float), \n",
    "                                     pos= torch.tensor(mesh.vertices, dtype=torch.float),\n",
    "                                     face = torch.tensor(mesh.faces, dtype=torch.long).t(),\n",
    "                                     edge_attr = torch.tensor(edge_attr, dtype=torch.float),\n",
    "                                     edge_index= torch.tensor(mesh.edges, dtype=torch.long).t().contiguous())\n",
    "    return data\n",
    "\n",
    "\n",
    "def process_mesh(path, suffix=\"\", model=None, out_dir=None, take_from_fld=True, prefields=None,\n",
    "                                  norm_field=False,  **kwargs):\n",
    "    FLD_PATH = path\n",
    "    if take_from_fld:\n",
    "        mesh, fields = load_fld(FLD_PATH)\n",
    "    else:\n",
    "        mesh, fields = load_predicted(FLD_PATH)\n",
    "    \n",
    "    if out_dir is None:\n",
    "        out_dir = os.path.join(*FLD_PATH.split(\"/\")[:-2], 'output')\n",
    "        \n",
    "    if model is not None:\n",
    "        if suffix == \"\":\n",
    "            suffix = '_predicted'\n",
    "        if take_from_fld:\n",
    "            data_instance = make_data_instance_from_stl(path)\n",
    "        else:\n",
    "            data_instance = make_data_instance_from_ply(path)\n",
    "        fields = model(data_instance.to('cuda:0')).cpu().detach().numpy()\n",
    "        \n",
    "    if prefields is not None:\n",
    "        fields = prefields\n",
    "\n",
    "    if norm_field:\n",
    "        fields = (fields - np.mean(fields[:, 0])) / np.std(fields[:, 0])\n",
    "        \n",
    "    image = visualize(mesh.vertices, mesh.faces, fields, **kwargs)\n",
    "\n",
    "    image_filename = os.path.join(out_dir, FLD_PATH.split(\"/\")[-1][:-4]) + suffix + \".png\" \n",
    "    imageio.imwrite(image_filename, image.astype(np.uint8))\n",
    "    \n",
    "def process_dir(path):\n",
    "    files = os.listdir(path)\n",
    "    for name in files:\n",
    "        process_mesh(os.path.join(path, name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n"
     ]
    }
   ],
   "source": [
    "model = SplineCNN8Residuals(3)\n",
    "model.load_state_dict(torch.load(\"Expirements/Networks15/normilized_full_latest.nn\"))\n",
    "model.to('cuda:0')\n",
    "model = model.eval()\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_mesh_opt(path, out, take_each=3, baseline=False, **kwargs):\n",
    "    if not os.path.exists(out):\n",
    "        os.makedirs(out)\n",
    "    \n",
    "    path = os.path.join(path, 'meshes')\n",
    "    files = [os.path.join(path, name) for name in filter(lambda x: x[0] == '0' and x[-3:] == \"ply\", os.listdir(path))]\n",
    "    for idx in [0] + list(range(5, len(files), take_each)): #(list(range(0, 30, take_each)) +\n",
    "        inp_path = files[idx]\n",
    "        data_instance = make_data_instance_from_ply(inp_path)\n",
    "        \n",
    "        if baseline:\n",
    "            if idx == 0:\n",
    "                edge_attrs = data_instance.edge_attr\n",
    "                continue\n",
    "            else:\n",
    "                data_instance.edge_attr = edge_attrs\n",
    "\n",
    "        fields = model(data_instance.to('cuda:0')).cpu().detach().numpy()\n",
    "        \n",
    "#         process_mesh(inp_path, suffix='_intr', prefields=fields,\n",
    "#                      out_dir=out, norm_field=True, **kwargs,\n",
    "#                      azimuth=240, elevation=5, take_from_fld=False)\n",
    "#         process_mesh(inp_path, suffix='_angl', prefields=fields,\n",
    "#                      out_dir=out, norm_field=True, **kwargs,\n",
    "#                      take_from_fld=False)\n",
    "#         process_mesh(inp_path, suffix='_pery', prefields=fields,\n",
    "#                      out_dir=out, norm_field=True, **kwargs,\n",
    "#                      azimuth=-270, elevation=90, take_from_fld=False)\n",
    "        process_mesh(inp_path, suffix='_perz', prefields=fields,\n",
    "                     out_dir=out, norm_field=True, **kwargs,\n",
    "                     azimuth=-270, elevation=0, take_from_fld=False)\n",
    "#         process_mesh(inp_path, suffix='_perx', prefields=fields,\n",
    "#                      out_dir=out, norm_field=True, **kwargs,\n",
    "#                      azimuth=180, elevation=0, take_from_fld=False)\n",
    "#         process_mesh(inp_path, suffix='_spoiler2', prefields=fields,\n",
    "#                      out_dir=out, norm_field=True,\n",
    "#                      azimuth=-45, elevation=0, take_from_fld=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for idx in [535]:\n",
    "#     visualize_mesh_opt('Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/%dminus/' % idx,\n",
    "#                        'Expirements/Visualizations/Paper/OptimizationDifferent/%dFreeFormMinus/' % idx, \n",
    "#                         baseline=True, take_each=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for idx in [69]:\n",
    "#     visualize_mesh_opt('Expirements/OptimizationPaper/AfterMeeting/UmetamiDrag2/%04d/' % idx,\n",
    "#                        'Expirements/Visualizations/Paper/OptimizationDifferent/%dUmetami/' % idx, \n",
    "#                         baseline=True, take_each=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/69minus/meshes'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-14-a7460622901c>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0midx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m535\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m69\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m32\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m162\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m61\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m     visualize_mesh_opt('Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/%dminus/' % idx,\n\u001b[0;32m----> 3\u001b[0;31m                        'Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/%dminus/frames' % idx, take_each=1, baseline=True)\n\u001b[0m",
      "\u001b[0;32m<ipython-input-11-c1197ae9af52>\u001b[0m in \u001b[0;36mvisualize_mesh_opt\u001b[0;34m(path, out, take_each, baseline, **kwargs)\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mpath\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'meshes'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m     \u001b[0mfiles\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfilter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'0'\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"ply\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlistdir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0midx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfiles\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtake_each\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m#(list(range(0, 30, take_each)) +\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m         \u001b[0minp_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfiles\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/69minus/meshes'"
     ]
    }
   ],
   "source": [
    "for idx in [535, 69, 32, 162, 61]:\n",
    "    visualize_mesh_opt('Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/%dminus/' % idx,\n",
    "                       'Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/%dminus/frames' % idx, take_each=1, baseline=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx in [175]:\n",
    "    visualize_mesh_opt('Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/%03d/' % idx,\n",
    "                       'Expirements/Visualizations/Paper/HighRes/FreeForm%04d/' % idx, take_each=19)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx in [175]:\n",
    "    visualize_mesh_opt('Expirements/OptimizationPaper/AfterMeeting/DeepSDFDrag//%03d/' % idx,\n",
    "                       'Expirements/Visualizations/Paper/HighRes/DeepSDF%04d/' % idx, take_each=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fields = np.load(a.replace('meshes', 'predictions').replace('ply', 'npy'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Expirements/OptimizationPaper/AfterMeeting/FreeformDrag/175/meshes/00003.ply'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plya"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_mesh_opt('Expirements/OptimizationPaper/AfterMeeting/DeepSDFDrag/175/',\n",
    "                   'Expirements/Visualizations/Paper/OptimizationDifferent/175SpoilerDisappear')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hotmap Visualizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dir = \"Expirements/Visualizations/Paper/PredictionComparison/afmhotNormFull_1\"\n",
    "colormap = cm.afmhot\n",
    "\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0001_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, colormap=colormap)\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0001_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, model=model, colormap=colormap)\n",
    "\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0002_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, colormap=cm.afmhot)\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0002_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, model=model, colormap=colormap)\n",
    "\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0003_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True)\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0003_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, model=model, colormap=colormap)\n",
    "\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0004_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, colormap=colormap)\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0004_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, model=model, colormap=colormap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dir = \"Expirements/Visualizations/Paper/PredictionComparison/hotNormFull_1\"\n",
    "colormap = cm.hot\n",
    "\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0001_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, colormap=colormap)\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0001_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, model=model, colormap=colormap)\n",
    "\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0002_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, colormap=cm.afmhot)\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0002_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, model=model, colormap=colormap)\n",
    "\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0003_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True)\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0003_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, model=model, colormap=colormap)\n",
    "\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0004_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, colormap=colormap)\n",
    "process_mesh('/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/0004_0015.fld', \n",
    "             out_dir=out_dir, norm_field=True, model=model, colormap=colormap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dir = \"Expirements/Visualizations/Paper/HighRes\"\n",
    "colormap = cm.jet\n",
    "\n",
    "for idx in [411]:\n",
    "    inp_path = '/cvlabdata2/home/artem/Data/cars_remeshed_dsdf/outputs/fld/%04d_0015.fld' % idx\n",
    "    if os.path.exists(inp_path):\n",
    "        process_mesh(inp_path, suffix='_spoilerHR_-120_10',\n",
    "                     out_dir=out_dir, norm_field=True,\n",
    "                     azimuth=-120, elevation=10)\n",
    "#                      out_dir=out_dir, norm_field=True, \n",
    "#                      azimuth=240, elevation=5, img_resolution=600)\n",
    "#         process_mesh(inp_path, suffix='_angl',\n",
    "#                      out_dir=out_dir, norm_field=True, img_resolution=600)\n",
    "#         process_mesh(inp_path, suffix='_pery',\n",
    "#                      out_dir=out_dir, norm_field=True,\n",
    "#                      azimuth=270, elevation=90,  img_resolution=600)\n",
    "#         process_mesh(inp_path, suffix='_perz', \n",
    "#                      out_dir=out_dir, norm_field=True,\n",
    "#                      azimuth=270, elevation=0, img_resolution=600)\n",
    "#         process_mesh(inp_path, suffix='_perx', \n",
    "#                      out_dir=out_dir, norm_field=True,\n",
    "#                      azimuth=180, elevation=0,  img_resolution=600)\n",
    "        \n",
    "#        process_mesh(inp_path, out_dir=out_dir, norm_field=True, model=model, colormap=colormap, img_resolution=600)\n",
    "    else:\n",
    "        print(\"No such file \", inp_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp_path = 'Expirements/OptimizationPaper/AfterMeeting/DeepSDFDrag/175/meshes/00039.ply'\n",
    "data_instance = make_data_instance_from_ply(inp_path)\n",
    "fields = model(data_instance.to('cuda:0')).cpu().detach().numpy()\n",
    "\n",
    "process_mesh (inp_path , prefields=fields,\n",
    "              out_dir=\"Expirements/Visualizations/Paper/OptimizationDifferent/175SpoilerDisappear\", \n",
    "              norm_field=True, suffix='_spoiler', azimuth=-30, elevation=0, take_from_fld=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display Distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh, fields = load_fld('/cvlabdata2/home/artem/Data/cars_refined/simulated/fld/0002_0005.fld')\n",
    "print( np.min(fields[:, 0]), np.max(fields[:, 0]) )\n",
    "norm_fields = (fields[:, 0] - np.mean(fields[:, 0])) / np.std(fields[:, 0])\n",
    "print(np.min(norm_fields), np.max(norm_fields))\n",
    "plt.hist(norm_fields, bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Draw Colormap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = plt.imshow(np.array([[-6, 6]]), cmap=\"jet\")\n",
    "img.set_visible(False)\n",
    "\n",
    "plt.colorbar(orientation=\"vertical\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pylab as pl\n",
    "import numpy as np\n",
    "\n",
    "a = np.array([[-1,1]])\n",
    "pl.figure(figsize=(1, 9))\n",
    "img = pl.imshow(a, cmap=\"jet\")\n",
    "pl.gca().set_visible(False)\n",
    "cb = pl.colorbar(orientation=\"vertical\", cax=pl.axes([0.1, 0.2, 0.4, 0.6]), ticks=[-0.8, 0, 0.8])\n",
    "lines = cb.ax.tick_params(size = 0, width = 5)\n",
    "pl.savefig(\"Expirements/Visualizations/Paper/PredictionComparison/jetColorMapOld/colorbar.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(lines[0].get_linewidths())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimisation Progress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = '/cvlabdata2/home/artem/DeepSDF/Expirements/OptimizationPaper/CleanedDataBadDrag/'\n",
    "\n",
    "for name in filter(lambda x: x[0] != '.' and x != 'DeepSDFDragFree', os.listdir(root)):\n",
    "    result = 0\n",
    "    num = 0\n",
    "    exp_dir = os.path.join(root, name)\n",
    "    for idx in filter(lambda x: x[0] != '.', os.listdir(exp_dir)):\n",
    "        for step_id in [0, 10, 20, 29]:\n",
    "            file_name = os.path.join(exp_dir, str(idx), 'meshes', str(step_id))\n",
    "            print(file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}