1 #!/usr/bin/env python -tt
   2 # encoding: utf-8
   3 #
   4 # Created by Holger Rapp on 2009-03-11.
   5 # HolgerRapp@gmx.net
   6 #
   7 
   8 import pycuda.driver as cuda
   9 import pycuda.compiler
  10 import pycuda.autoinit
  11 import numpy
  12 from math import pi,cos,sin
  13 
  14 _rotation_kernel_source = """
  15 texture<float, 2> tex;
  16 
  17 __global__ void copy_texture_kernel(
  18     const float resize_val, 
  19     const float alpha, 
  20     unsigned short oldiw, unsigned short oldih,
  21     unsigned short newiw, unsigned short newih,
  22     unsigned char* data) {
  23 
  24         // calculate pixel idx
  25         unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;
  26         unsigned int y = blockIdx.y * blockDim.y + threadIdx.y;
  27         
  28         // We might be outside the reachable pixels. Don't do anything
  29         if( (x >= newiw) || (y >= newih) )
  30             return;
  31         
  32         // calculate offset into destination array
  33         unsigned int didx = y * newiw + x;
  34         
  35         // calculate offset into source array (be aware of rotation and scaling)
  36         float xmiddle = (x-newiw/2.) / resize_val;
  37         float ymiddle = (y-newih/2.) / resize_val;
  38         float sx = ( xmiddle*cos(alpha)+ymiddle*sin(alpha) + oldiw/2.) ;
  39         float sy = ( -xmiddle*sin(alpha)+ymiddle*cos(alpha) + oldih/2.);
  40         
  41         if( (sx < 0) || (sx >= oldiw) || (sy < 0) || (sy >= oldih) ) { 
  42             data[didx] = 255; 
  43             return;
  44         }
  45 
  46         data[didx] = tex2D(tex, sx, sy);
  47     }
  48 """
  49 mod_copy_texture=pycuda.compiler.SourceModule( _rotation_kernel_source )
  50 
  51 copy_texture_func = mod_copy_texture.get_function("copy_texture_kernel")
  52 texref = mod_copy_texture.get_texref("tex")
  53 
  54 def rotate_image( a, resize = 1.5, angle = 20., interpolation = "linear", blocks = (16,16,1)  ):
  55     """
  56     Rotates the array. The new array has the new size and centers the
  57     picture in the middle.
  58     
  59     a             - array (2-dim)
  60     resize        - new_image w/old_image w
  61     angle         - degrees to rotate the image
  62     interpolation - "linear" or None
  63     blocks        - given to the kernel when run
  64 
  65 
  66     returns: a new array with dtype=uint8 containing the rotated image
  67     """
  68     angle = angle/180. *pi
  69     
  70     # Convert this image to float. Unsigned int texture gave 
  71     # strange results for me. This conversion is slow though :(
  72     a = a.astype("float32")
  73 
  74     # Calculate the dimensions of the new image
  75     calc_x = lambda (x,y): (x*a.shape[1]/2.*cos(angle)-y*a.shape[0]/2.*sin(angle))
  76     calc_y = lambda (x,y): (x*a.shape[1]/2.*sin(angle)+y*a.shape[0]/2.*cos(angle))
  77 
  78     xs = [ calc_x(p) for p in [ (-1.,-1.),(1.,-1.),(1.,1.),(-1.,1.) ] ]
  79     ys = [ calc_y(p) for p in [ (-1.,-1.),(1.,-1.),(1.,1.),(-1.,1.) ] ]
  80 
  81     new_image_dim = (
  82         int(numpy.ceil(max(ys)-min(ys))*resize),
  83         int(numpy.ceil(max(xs)-min(xs))*resize),
  84     )
  85     
  86     # Now generate the cuda texture
  87     cuda.matrix_to_texref(a, texref, order="C")
  88     
  89     # We could set the next if we wanted to address the image
  90     # in normalized coordinates ( 0 <= coordinate < 1.)
  91     # texref.set_flags(cuda.TRSF_NORMALIZED_COORDINATES)
  92     if interpolation == "linear":
  93         texref.set_filter_mode(cuda.filter_mode.LINEAR)
  94 
  95     # Calculate the gridsize. This is entirely given by the size of our image. 
  96     gridx = new_image_dim[0]/blocks[0] if \
  97             new_image_dim[0]%blocks[0]==1 else new_image_dim[0]/blocks[0] +1
  98     gridy = new_image_dim[1]/blocks[1] if \
  99             new_image_dim[1]%blocks[1]==0 else new_image_dim[1]/blocks[1] +1
 100 
 101     # Get the output image
 102     output = numpy.zeros(new_image_dim,dtype="uint8")
 103     
 104     # Call the kernel
 105     copy_texture_func(
 106         numpy.float32(resize), numpy.float32(angle),
 107         numpy.uint16(a.shape[1]), numpy.uint16(a.shape[0]),
 108         numpy.uint16(new_image_dim[1]), numpy.uint16(new_image_dim[0]),
 109             cuda.Out(output),texrefs=[texref],block=blocks,grid=(gridx,gridy))
 110     
 111     return output
 112 
 113 if __name__ == '__main__':
 114     import Image
 115     import sys
 116     
 117     def main( ):
 118         if len(sys.argv) != 2:
 119             print "You should really read the source...\n\nUsage: rotate.py <Imagename>\n"
 120             sys.exit(-1)
 121 
 122         # Open, convert to grayscale, convert to numpy array
 123         img = Image.open(sys.argv[1]).convert("L")
 124         i = numpy.fromstring(img.tostring(),dtype="uint8").reshape(img.size[1],img.size[0])
 125         
 126         # Rotate & convert back to PIL Image
 127         irot = rotate_image(i)
 128         rotimg = Image.fromarray(irot,mode="L")
 129 
 130         # Save and display
 131         rotimg.save("rotated.png")
 132         rotimg.show()
 133     
 134     main()