1 """
   2 3D display of Light Field images.
   3 Example images can be download from:
   4 http://graphics.stanford.edu/software/LFDisplay/LFDisplay-samples.zip
   5 
   6 Use:
   7 >> python LFview.py <path to light-field image>
   8 
   9 Prerequisites:
  10 - Python 2.7
  11 - Enthought Tool Suite (ETS)
  12 - PIL
  13 - Jinja
  14 
  15 Author: Amit Aides. amitibo at technion . ac . il
  16 """
  17 
  18 from __future__ import division
  19 
  20 from enthought.traits.api import HasTraits, Range, on_trait_change
  21 from enthought.traits.ui.api import View, Item
  22 from enthought.chaco.api import Plot, ArrayPlotData, gray
  23 from enthought.enable.component_editor import ComponentEditor
  24 
  25 import numpy as np
  26 import Image
  27 import argparse
  28 import os.path
  29 import math
  30 
  31 import pycuda.driver as cuda
  32 import pycuda.compiler
  33 import pycuda.autoinit
  34 
  35 from jinja2 import Template
  36 
  37 
  38 _kernel_tpl = Template("""
  39 {% if NCHANNELS == 3 %}
  40 texture<float4, 2> tex;
  41 {% else %}
  42 texture<float, 2> tex;
  43 {% endif %}
  44 
  45 __global__ void LFview_kernel(
  46     float x_angle,
  47     float y_angle,
  48     unsigned char* data
  49     )
  50 {
  51     //
  52     // calculate pixel idx
  53     //
  54     unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;
  55     unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;
  56 
  57     //
  58     // We might be outside the reachable pixels. Don't do anything
  59     //
  60     if( (x >= {{newiw}}) || (y >= {{newih}}) )
  61         return;
  62 
  63     //
  64     // calculate offset into destination array
  65     //
  66     unsigned int didx = (y * {{newiw}} + x) * {{NCHANNELS}};
  67     
  68     //
  69     // calculate offset into source array (be aware of rotation and scaling)
  70     //
  71     float sx = {{x_start}} + tan(x_angle)*{{x_ratio}} + {{x_step}}*x;
  72     float sy = {{y_start}} + tan(y_angle)*{{y_ratio}} + {{y_step}}*y;
  73 
  74     if( (sx < 0) || (sx >= {{oldiw}}) || (sy < 0) || (sy >= {{oldih}}) ) {
  75 
  76         {% for channel in range(NCHANNELS) %}
  77         data[didx+{{channel}}] = 0;
  78         {% endfor %}
  79 
  80         return;
  81     }
  82 
  83     {% if NCHANNELS == 3 %}
  84     float4 texval = tex2D(tex, sx, sy);
  85     data[didx] = texval.x;
  86     data[didx+1] = texval.y;
  87     data[didx+2] = texval.z;
  88     {% else %}
  89     data[didx] = tex2D(tex, sx, sy);
  90     {% endif %}
  91 }
  92 """)
  93 
  94 
  95 def ceil(x):
  96     return int(x + 0.5)
  97 
  98 
  99 class LFapplication(HasTraits):
 100 
 101     traits_view = View(
 102         Item('LF_img', editor=ComponentEditor(), show_label=False),
 103         Item('X_angle', label='Angle in the X axis'),
 104         Item('Y_angle', label='Angle in the Y axis'),
 105         resizable = True,
 106         title="LF Image"
 107         )
 108 
 109     def __init__(self, img_path):
 110         super(LFapplication, self).__init__()
 111 
 112         #
 113         # Load image data
 114         #
 115         base_path = os.path.splitext(img_path)[0]
 116         lenslet_path = base_path + '-lenslet.txt'
 117         optics_path = base_path + '-optics.txt'
 118 
 119         with open(lenslet_path, 'r') as f:
 120             tmp = eval(f.readline())
 121             x_offset, y_offset, right_dx, right_dy, down_dx, down_dy = \
 122               np.array(tmp, dtype=np.float32)
 123 
 124         with open(optics_path, 'r') as f:
 125             for line in f:
 126                 name, val = line.strip().split()
 127                 try:
 128                     setattr(self, name, np.float32(val))
 129                 except:
 130                     pass
 131 
 132         max_angle = math.atan(self.pitch/2/self.flen)
 133 
 134         #
 135         # Prepare image
 136         #
 137         im_pil = Image.open(img_path)
 138         if im_pil.mode == 'RGB':
 139             self.NCHANNELS = 3
 140             w, h = im_pil.size
 141             im = np.zeros((h, w, 4), dtype=np.float32)
 142             im[:, :, :3] = np.array(im_pil).astype(np.float32)
 143             self.LF_dim = (ceil(h/down_dy), ceil(w/right_dx), 3)
 144         else:
 145             self.NCHANNELS = 1
 146             im = np.array(im_pil.getdata()).reshape(im_pil.size[::-1]).astype(np.float32)
 147             h, w = im.shape
 148             self.LF_dim = (ceil(h/down_dy), ceil(w/right_dx))
 149 
 150         x_start = x_offset - int(x_offset / right_dx) * right_dx
 151         y_start = y_offset - int(y_offset / down_dy) * down_dy
 152         x_ratio = self.flen * right_dx / self.pitch
 153         y_ratio = self.flen * down_dy / self.pitch
 154 
 155         #
 156         # Generate the cuda kernel
 157         #
 158         mod_LFview = pycuda.compiler.SourceModule(
 159             _kernel_tpl.render(
 160                 newiw=self.LF_dim[1],
 161                 newih=self.LF_dim[0],
 162                 oldiw=w,
 163                 oldih=h,
 164                 x_start=x_start,
 165                 y_start=y_start,
 166                 x_ratio=x_ratio,
 167                 y_ratio=y_ratio,
 168                 x_step=right_dx,
 169                 y_step=down_dy,
 170                 NCHANNELS=self.NCHANNELS
 171                 )
 172             )
 173         
 174         self.LFview_func = mod_LFview.get_function("LFview_kernel")
 175         self.texref = mod_LFview.get_texref("tex")
 176         
 177         #
 178         # Now generate the cuda texture
 179         #
 180         if self.NCHANNELS == 3:
 181             cuda.bind_array_to_texref(
 182                 cuda.make_multichannel_2d_array(im, order="C"),
 183                 self.texref
 184                 )
 185         else:
 186             cuda.matrix_to_texref(im, self.texref, order="C")
 187             
 188         #
 189         # We could set the next if we wanted to address the image
 190         # in normalized coordinates ( 0 <= coordinate < 1.)
 191         # texref.set_flags(cuda.TRSF_NORMALIZED_COORDINATES)
 192         #
 193         self.texref.set_filter_mode(cuda.filter_mode.LINEAR)
 194 
 195         #
 196         # Prepare the traits
 197         #
 198         self.add_trait('X_angle', Range(-max_angle, max_angle, 0.0))
 199         self.add_trait('Y_angle', Range(-max_angle, max_angle, 0.0))
 200         
 201         self.plotdata = ArrayPlotData(LF_img=self.sampleLF())
 202         self.LF_img = Plot(self.plotdata)
 203         if self.NCHANNELS == 3:
 204             self.LF_img.img_plot("LF_img")
 205         else:
 206             self.LF_img.img_plot("LF_img", colormap=gray)
 207 
 208     def sampleLF(self):
 209         #
 210         # Get the output image
 211         #
 212         output = np.zeros(self.LF_dim, dtype=np.uint8)
 213         
 214         #
 215         # Calculate the gridsize. This is entirely given by the size of our image. 
 216         #
 217         blocks = (16, 16, 1)
 218         gridx = ceil(self.LF_dim[1]/blocks[1])
 219         gridy = ceil(self.LF_dim[0]/blocks[0])
 220         grid = (gridx, gridy)
 221 
 222         #
 223         # Call the kernel
 224         #
 225         self.LFview_func(
 226             np.float32(self.X_angle),
 227             np.float32(self.Y_angle),
 228             cuda.Out(output),
 229             texrefs=[self.texref],
 230             block=blocks,
 231             grid=grid
 232             )
 233 
 234         return output
 235 
 236     @on_trait_change('X_angle, Y_angle')        
 237     def updateImge(self):
 238         self.plotdata.set_data('LF_img', self.sampleLF())
 239         
 240         
 241 def main(img_path):
 242     """Main function"""
 243 
 244     app = LFapplication(img_path)
 245     app.configure_traits()
 246     
 247     
 248 if __name__ == '__main__':
 249     parser = argparse.ArgumentParser(description='View an LF image')
 250     parser.add_argument('img_path', type=str, help='Path to LF image')
 251     args = parser.parse_args()
 252 
 253     main(args.img_path)