make_grid_np -> make_proj_grid, simplify fn

AllanYangZhou · May 1, 2018 · d1f375d · d1f375d
1 parent 3920ac0
commit d1f375d
Show file tree

Hide file tree

Showing 4 changed files with 131 additions and 116 deletions.
diff --git a/PytorchSphProjectionExample.ipynb b/PytorchSphProjectionExample.ipynb
diff --git a/StyleTransfer-PyTorch-Baseline.ipynb b/StyleTransfer-PyTorch-Baseline.ipynb
@@ -26,7 +26,9 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "metadata": {},
+   "metadata": {
+    "collapsed": false
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -56,8 +58,6 @@
     "#matplotlib.use('TkAgg')\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "from projection_utils import make_grid_np\n",
-    "\n",
     "SQUEEZENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)\n",
     "SQUEEZENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)\n",
     "%matplotlib inline"
@@ -73,7 +73,9 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "def preprocess(img, size=512):\n",
@@ -129,7 +131,9 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "#dtype = torch.FloatTensor\n",
@@ -141,7 +145,9 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "metadata": {},
+   "metadata": {
+    "collapsed": false
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -215,7 +221,9 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "def content_loss(content_weight, content_current, content_original):\n",
@@ -261,7 +269,9 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "def gram_matrix(features, normalize=True):\n",
@@ -302,7 +312,9 @@
   {
    "cell_type": "code",
    "execution_count": 7,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "# Now put it together in the style_loss function...\n",
@@ -353,7 +365,9 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "def tv_loss(img, tv_weight):\n",
@@ -384,18 +398,21 @@
   {
    "cell_type": "code",
    "execution_count": 9,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "# TODO: randomly theta,phi unfixed.\n",
-    "coords = make_grid_np(np.pi/3, 0, 200, (512, 256, 3))\n",
+    "coords = make_proj_grid(np.pi/3, 0, 200, (512, 256, 3))\n",
     "coords_var = Variable(torch.from_numpy(coords.astype(np.float32))[None], requires_grad=False).cuda()"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 20,
    "metadata": {
+    "collapsed": true,
     "scrolled": false
    },
    "outputs": [],
@@ -497,7 +514,9 @@
   {
    "cell_type": "code",
    "execution_count": 21,
-   "metadata": {},
+   "metadata": {
+    "collapsed": false
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -621,7 +640,9 @@
   {
    "cell_type": "code",
    "execution_count": 22,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "im = deprocess(grad.data.cpu().clone())\n",
@@ -631,7 +652,9 @@
   {
    "cell_type": "code",
    "execution_count": 23,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "im.save('result.jpg')"
@@ -640,7 +663,9 @@
   {
    "cell_type": "code",
    "execution_count": 28,
-   "metadata": {},
+   "metadata": {
+    "collapsed": false
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -764,7 +789,9 @@
   {
    "cell_type": "code",
    "execution_count": 29,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "im = deprocess(grad.data.cpu().clone())\n",
@@ -774,7 +801,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "collapsed": true
+   },
    "outputs": [],
    "source": [
     "# Composition VII + Tubingen\n",
@@ -796,7 +825,9 @@
   {
    "cell_type": "code",
    "execution_count": 58,
-   "metadata": {},
+   "metadata": {
+    "collapsed": false
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -888,7 +919,9 @@
   {
    "cell_type": "code",
    "execution_count": 59,
-   "metadata": {},
+   "metadata": {
+    "collapsed": false
+   },
    "outputs": [
     {
      "name": "stderr",

diff --git a/main.py b/main.py
@@ -8,7 +8,7 @@
 import matplotlib
 import matplotlib.pyplot as plt
 
-from projection_utils import make_grid_np
+from projection_utils import make_proj_grid
 from style_utils import extract_features, content_loss, style_loss, tv_loss, gram_matrix
 from image_utils import *
 
@@ -84,8 +84,7 @@ def style_transfer(cnn, content_image, style_image, image_size, style_size, cont
 
         for i in range(num_proj):
             theta, phi = np.random.uniform(0, 2 * np.pi), np.random.uniform(0, 2 * np.pi)
-            grid_size = (image_size[0], image_size[1], 3)
-            coords = make_grid_np(theta, phi, 500, grid_size )
+            coords = make_proj_grid(theta, phi, 500)
             coords_var = Variable(torch.from_numpy(coords.astype(np.float32))[None].type(dtype), requires_grad=False)
 
             proj_content_img_var = F.grid_sample(content_img_var, coords_var)

diff --git a/projection_utils.py b/projection_utils.py
@@ -1,4 +1,4 @@
-# Adapted from: https://github.com/bingsyslab/360projection/blob/master/equirectangular.py
+# Adapted from: https://github.com/bingsyslab/360projection/
 
 import numpy as np
 
@@ -12,21 +12,17 @@ def yrotation(th):
     s = np.sin(th)
     return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])
 
-def make_grid_np(theta0, phi0, width, img_shape, fov_h=np.pi/2, fov_v=np.pi/2):
+def make_proj_grid(theta0, phi0, width, fov_h=np.pi/2, fov_v=np.pi/2):
     """
     theta0 is pitch
     phi0 is yaw
-    render view at (pitch, yaw) with fov_h by fov_v
     width is the number of horizontal pixels in the view
 
-    Returns a (height, width, 2) matrix containing the
-    pixels to sample from (in the source img)
+    Returns a (height, width, 2) matrix which you can use
     to project from sphere->perspective.
     """
     m = np.dot(yrotation(phi0), xrotation(theta0))
 
-    (base_height, base_width, _) = img_shape
-
     height = int(width * np.tan(fov_v / 2) / np.tan(fov_h / 2))
 
     DI = np.ones((height * width, 3), np.int)
@@ -47,13 +43,10 @@ def make_grid_np(theta0, phi0, width, img_shape, fov_h=np.pi/2, fov_v=np.pi/2):
     theta = np.pi / 2 - np.arctan2(v[:, 1], diag)
     phi = np.arctan2(v[:, 0], v[:, 2]) + np.pi
 
-    ey = theta * base_height / np.pi
-    ex = phi * base_width / (2 * np.pi)
-
-    ex[ex >= base_width] = base_width - 1
-    ey[ey >= base_height] = base_height - 1
+    ey = theta / np.pi
+    ex = phi / (2 * np.pi)
 
-    # Need to scale grid coordinates between [-1, +1] for use in grid_sample
-    ey_grid = ey.reshape(height, width) / (base_height / 2.) - 1
-    ex_grid = ex.reshape(height, width) / (base_width / 2.) - 1
+    # Scale grid coordinates between [-1, +1]
+    ey_grid = ey.reshape(height, width) * 2. - 1
+    ex_grid = ex.reshape(height, width) * 2. - 1
     return np.stack([ex_grid, ey_grid], axis=-1)