

import Image
from numpy import *
from scipy.ndimage import geometric_transform
from scipy.misc import imsave
import time
from mpi4py import MPI


comm = MPI.COMM_WORLD  # comunicador entre dos procesadores
rank = comm.rank     # id procesador actual
size = comm.size     # cantidad de procesadores a usar

def polar_warp(im,sz,r):
    #""" Warp an image to a square polar version. """
    h,w = im.shape[:2]

    def polar_2_rect(xy):
        #""" Convert polar coordinates to coordinates in
        #	the original image for geometric_transform. """
        x,y = float(xy[1]),float(xy[0])
        theta = 0.5*arctan2(x-sz/2,y-sz/2)
        R = sqrt( (x-sz/2)*(x-sz/2) + (y-sz/2)*(y-sz/2) ) - r
        return (2*h*R/(sz-2*r),w/2+theta*w/pi+pi/2)

    # if color image, warp each channel, otherwise there's just one
    if len(im.shape)==3:
        warped = zeros((sz,sz,3))
        for i in range(3):
            warped[:,:,i] = geometric_transform(im[:,:,i],polar_2_rect,output_shape=(sz,sz),mode='nearest')
    else:
        warped = geometric_transform(im,polar_2_rect,output_shape=(sz,sz),mode='nearest')
    return warped


if __name__ == "__main__":

    # load image from http://www.filipecalegario.com/comppho/3assignment/
    starting_point=time.time()
    im = array(Image.open('panoramica.jpg'))
    h,w = im.shape[:2]
    # flip vertical to make math easier for warp
    im = flipud(im)
    # warp image
    wim = uint8(polar_warp(im,h,5))
    # save
    imsave('polar.png',wim)
    elapsed_time=time.time()-starting_point
    if rank==0:
        print ""
        print "Serial Time [seconds]: " + str(elapsed_time)