Stereogram generator

stereogram.py

#!/usr/bin/env python

from PIL import Image
import numpy as np
from sys import argv,exit
from random import randint
from time import time
from math import ceil
try:
        from weave import inline,converters
except:
        from scipy.weave import inline,converters

def displacement_map(fname):
        return displacement_map_img(Image.open(fname))

def displacement_map_img(img):
        if img.mode != 'L':
                img = img.convert('L')
        mat = np.asarray(img)
        mx,mn = np.max(mat),np.min(mat)
        mat = (120+(mat-mn)*40.0/(mx-mn)).astype(np.uint16)
        return mat

def pattern(fname,nr,nc):
        img = Image.open(fname)
        if img.mode != 'RGB':
                img = img.convert('RGB')
        patt = np.asarray(img)[:nr,:nc]
        pr,pc,_ = patt.shape
        patt = np.tile(patt,(ceil(nr/float(pr)),ceil(nc/float(pc)),1))
        return patt[:nr,:nc]

def pattern_asg(disp_map,patt):
        asg = np.zeros(disp_map.shape+(3,),dtype=np.int16)-1
        num_rows,num_cols = disp_map.shape
        code = """
        int c,cd,colR,colG,colB;
        for (int r=0; r<num_rows; ++r){
                c = cd = 0;
                while (c < num_cols){
                        colR = patt(r,c,0);
                        colG = patt(r,c,1);
                        colB = patt(r,c,2);
                        while (cd < num_cols){
                                asg(r,cd,0) = colR;
                                asg(r,cd,1) = colG;
                                asg(r,cd,2) = colB;
                                cd += disp_map(r,cd);
                        }
                        while (c < num_cols && asg(r,c,0) != -1)
                                ++c;
                        cd = c;
                }
        }"""
        inline(code,['num_rows','num_cols','asg','disp_map','patt'],
                type_converters=converters.blitz)
        return asg.astype(np.uint8)

def random_dot_asg(disp_map):
        asg = np.zeros(disp_map.shape+(3,),dtype=np.uint8)
        num_rows,num_cols = disp_map.shape
        code = """
        int c,cd,colR,colG,colB;
        for (int r=0; r<num_rows; ++r){
                c = cd = 0;
                while (c < num_cols){
                        colR = rand()%255+1;
                        colG = rand()%255+1;
                        colB = rand()%255+1;
                        while (cd < num_cols){
                                asg(r,cd,0) = colR;
                                asg(r,cd,1) = colG;
                                asg(r,cd,2) = colB;
                                cd += disp_map(r,cd);
                        }
                        while (c < num_cols && asg(r,c,0))
                                ++c;
                        cd = c;
                }
        }"""
        inline(code,['num_rows','num_cols','asg','disp_map'],
                type_converters=converters.blitz)
        return asg

def write_asg(asg,fname):
        Image.fromarray(asg,'RGB').save(fname)

if __name__ == "__main__":
        if not len(argv) in [3,4]:
                exit("Usage: %s depthmap.img outfile.img [pattern.img]"%argv[0])

        dmap = displacement_map(argv[1])
        if len(argv) == 3:
                asg = random_dot_asg(dmap)
        else:
                asg = pattern_asg(dmap,pattern(argv[3],*dmap.shape))
        write_asg(asg,argv[2])