optical_illusion.py

# -*- coding: utf-8 -*-

'''
錯視図の作成
'''

import numpy as np
import time
import argparse
import matplotlib.pylab as plt


def projection(x, y, f, theta, cz):
    # imageから机上への投影
    xdash = x * cz / (f * np.cos(theta) - y * np.sin(theta))
    ydash = cz * (f*np.sin(theta) + y*np.cos(theta)) / \
                 (f*np.cos(theta) - y*np.sin(theta))
    return np.array([xdash, ydash]).T

def reprojection(xdash, ydash, f, theta, cz):
    # 机上からimageへの投影
    y = f * (ydash*np.cos(theta) - cz*np.sin(theta))/ \
            (ydash*np.sin(theta) + cz*np.cos(theta))
    x = xdash / cz * (f * np.cos(theta) - y * np.sin(theta))
    return np.array([x, y]).T

def projection_range(img, sensorhight, f, theta, z):
    # imageの投影先範囲を計算
    (xlength, ylength, dim) = img.shape
    xmin = -sensorhight*xlength/ylength/3
    xmax =  sensorhight*xlength/ylength/3
    ymin = -sensorhight/2
    ymax =  sensorhight/2
    
    xdashmin = projection(xmin, ymax, f, theta, z)
    xdashmin = xdashmin[0]
    xdashmax = projection(xmax, ymax, f, theta, z)
    xdashmax = xdashmax[0]
    ydashmin = projection(0, ymin, f, theta, z)
    ydashmin = ydashmin[1]
    ydashmax = projection(0, ymax, f, theta, z)
    ydashmax = ydashmax[1]

    return (xdashmin, xdashmax, ydashmin, ydashmax)

def gen_blanckimage(xmin, xmax, ymin, ymax, notch = 0.01):
    # 空白イメージを作成し、イメージと対応する座標を返却
    xpick = np.arange(xmin, xmax, notch)
    ypick = np.arange(ymin, ymax, notch)
    outimg = np.ones([ypick.size, xpick.size, 4], dtype=np.float32)
    return (outimg, xpick, ypick)

def pixcel_copy(inimg, outimg, fromx, fromy, tox, toy):
    (ymax, xmax, dim) = inimg.shape
    for i in range(toy.size):
        # set color
        if (0 <= fromx[i] < xmax) & (0 <= fromy[i] < ymax):
            outimg[toy[i], tox[i], :] = inimg[fromy[i], fromx[i], :]
    return outimg

def rasterize(img, f, theta, z, sensorhight, xpick, ypick, outimg):
    # 画素の値を抽出
    udimg = np.flipud(img)
    (ymax, xmax, dim) = udimg.shape
    for ix in range(xpick.size):
        xyp = reprojection(xpick[ix], ypick, f, theta, z).T
        xp = xyp[0]
        yp = xyp[1]
        
        # 空間座標からピクセル座標への変換
        fromx = np.round(xp * ymax / sensorhight) + xmax / 2
        fromy = np.round(yp * ymax / sensorhight) + ymax / 2
        tox = np.repeat(ix, ypick.shape[0])
        toy = np.arange(ypick.size)
        outimg = pixcel_copy(udimg, outimg, fromx, fromy, tox, toy)
    return np.flipud(outimg)

def gen_image(img, f, theta, z, sensorhight, notch=0.01):
    # 机上への投影イメージを作成
    (xmin, xmax, ymin, ymax) = projection_range(img, sensorhight, f, theta, z)
    (outimg, xpick, ypick) = gen_blanckimage(xmin, xmax, ymin, ymax, notch)
    outimg = rasterize(img, f, theta, z, sensorhight, xpick, ypick, outimg)
    return outimg

def test():
    start = time.time()
    
    img = plt.imread('image_noground.png')
    f=35
    theta=np.deg2rad(50)
    z=7
    sensorhight = 18
    outimg = gen_image(img, f, theta, z, sensorhight, notch=0.01)
    
    elapsed_time = time.time() - start
    print( ("elapsed_time:{0}".format(elapsed_time)) + "[sec]" )
    
    plt.imshow(outimg)

if __name__ == '__main__':
    # pars args
    parser = argparse.ArgumentParser()
    parser.add_argument("input", help="input image file name")
    parser.add_argument("focus", type=float, help="focus of camera")
    parser.add_argument("theta", type=float,
                        help="angle of camera from perpendicular")
    parser.add_argument("z", type=float, help="hight of camera")
    parser.add_argument("sensorhight", type=float, help="sensor hight of camera")
    parser.add_argument("output", help="output image file name")
    parser.add_argument("-n", "--notch", type=float, 
                        help="notch of output image(default 0.01)")
    args = parser.parse_args()
    
    # set parameter
    img = plt.imread(args.input)
    f = args.focus
    theta = np.deg2rad(args.theta)
    z = args.z
    sensorhight = args.sensorhight
    
    notch = 0.01
    if args.notch:
        notch = args.notch
    
    # gen image
    outimg = gen_image(img, f, theta, z, sensorhight, notch)
    plt.imsave(args.output, outimg)