#!/usr/bin/env python3 import os,sys,math import numpy as np import cv2 import matplotlib matplotlib.use('tkagg') #CV2 causes some gear grinding with matplotlib's version of Qt. This fixes it via magic bullshit. import matplotlib.pyplot as plt import scipy import scipy.interpolate as interp import threading #this is actually useless for multithreaded execution! import multiprocessing as mp from multiprocessing import shared_memory cos=np.cos sin=np.sin tan=np.tan atan=np.arctan pi = np.pi sqrt = np.sqrt exp = np.exp """ Script to stitch multiple output images into a full sky hdr image. Scroll down to the bottom for the main routine. """ def load_img(fname): """ Data plumbing to fix all the stupid things that opencv does in terms of color plane order, etc output will be an RGBA image (Nx,Ny,4), float32, scaled from 0 to 1, and """ try: img = cv2.imread(fname,cv2.IMREAD_UNCHANGED) except: print("load_img: error: could not open {} as an image with cv2".format(fname)) return None Nx = img.shape[0] Ny = img.shape[1] imgout = np.zeros((Nx,Ny,4),dtype=np.float32) if(len(img.shape)<=2): #grayscale? imgout[:,:,0] = img[:,:] imgout[:,:,1] = img[:,:] imgout[:,:,2] = img[:,:] imgout[:,:,3] = 255.0 imgout = imgout/255.0 imgout = np.clip(imgout,0.0,1.0) imgout = np.transpose(imgout,(1,0,2)) else: if(img.shape[2]==3): #BGR image imgout[:,:,0] = img[:,:,2] imgout[:,:,1] = img[:,:,1] imgout[:,:,2] = img[:,:,0] imgout[:,:,3] = 255.0 imgout = imgout/255.0 imgout = np.clip(imgout,0.0,1.0) imgout = np.transpose(imgout,(1,0,2)) elif(img.shape[2]>=4): #BGRA image imgout[:,:,0] = img[:,:,2] imgout[:,:,1] = img[:,:,1] imgout[:,:,2] = img[:,:,0] imgout[:,:,3] = img[:,:,3] imgout = imgout/255.0 imgout = np.clip(imgout,0.0,1.0) imgout = np.transpose(imgout,(1,0,2)) return imgout def save_img(fname,img): """ takes a floating point (0,1) normalized image greyscale: (Nx,Ny) RGB: (Nx,Ny,3) RGBA: (Nx,Ny,4) saves it to fname via cv2 """ Ny = img.shape[0] Nx = img.shape[1] if(len(img.shape)==2): Ncp = 1 img = img.copy() img = img.transpose() else: Ncp = img.shape[2] img = img.copy() img = img.transpose(1,0,2) im2 = np.zeros((Nx,Ny,4),dtype=np.float32) im3 = np.zeros((Nx,Ny,4),dtype=np.uint8) if(Ncp==1): im2[:,:,0] = img[:,:] im2[:,:,1] = img[:,:] im2[:,:,2] = img[:,:] im2[:,:,3] = 1.0 elif(Ncp==3): im2[:,:,0] = img[:,:,2] im2[:,:,1] = img[:,:,1] im2[:,:,2] = img[:,:,0] im2[:,:,3] = 1.0 elif(Ncp>=4): im2[:,:,0] = img[:,:,2] im2[:,:,1] = img[:,:,1] im2[:,:,2] = img[:,:,0] im2[:,:,3] = img[:,:,3] else: print("save_img: no way to handle Ncp={}".format(Ncp)) im2 = im2*255.0 im2 = np.clip(im2,0.0,255.0) im3 = np.uint8(im2) try: cv2.imwrite(fname,im3) except: print("save_img: error: could not write to file {} fname".format(fname)) return def test_img_loadsave(): img = load_img("./testimg.png") print(img.shape) plt.imshow(img.transpose(1,0,2)) plt.show() save_img("./testout.png",img) img = np.random.rand(255,128) save_img("./testout2.png",img) img = np.random.rand(255,128,3) save_img("./testout3.png",img) ### ### Some vector math ### def hodge_dual(v): """ Hodge dual of a vector. The hodge dual is an operation in differentail forms that (in 3d) returns a bivector for a vector, the bivector being related to the generator of a rotation matrix. """ v = np.array(v) M = np.zeros((3,3)) M[0,1] = v[2] M[1,2] = v[0] M[2,0] = v[1] M[1,0] = -v[2] M[2,1] = -v[0] M[0,2] = -v[1] M = -M return M def rotmat_axisangle(axis,angle): """ rotmat_axisangle(axis,angle) axis - the axis of rotation angle - the angle [rad] to rotate about the axis returns: 3x3 rotation matrix """ axis = np.array(axis) E = np.eye(3) H = hodge_dual(axis) H2 = np.matmul(H,H) R = E + H*sin(angle) + H2*(1.0-cos(angle)) return R def rotx(angle): R = np.array([[1,0,0],[0,cos(angle),-sin(angle)],[0,sin(angle),cos(angle)]]) return R def roty(angle): R = np.array([[cos(angle),0,sin(angle)],[0,1,0],[-sin(angle),0,cos(angle)]]) return R def rotz(angle): R = np.array([[cos(angle),-sin(angle),0],[sin(angle),cos(angle),0],[0,0,1]]) return R def rot_eulerxyz(x,y,z): Rx = rotx(x) Ry = roty(y) Rz = rotz(z) R = np.matmul(Rz,np.matmul(Ry,Rx)) return R def rot_cameraxyz(x,y,z): R1 = rotx(-90*pi/180) R2 = rot_eulerxyz(x,y,z) R = np.matmul(R2,R1) return R def rot_eulerzyx(x,y,z): Rx = rotx(x) Ry = roty(y) Rz = rotz(z) R = np.matmul(Rx,np.matmul(Ry,Rz)) return R def extrinsic(R,t): R = np.array(R).copy() t = np.array(t).copy() t = np.reshape(t,[3,1]) Rinv = np.linalg.inv(R) c = -np.matmul(Rinv,t) Ex = np.zeros([4,4]) Ex[0:3,0:3] = Rinv[:,:] Ex[0:3,3] = c.flatten() Ex[3,3] = 1.0 return Ex def intrinsic_eqasp(fov): """ Equal aspect ratio intrinsic matrix """ In = np.zeros([4,4]) alpha = fov/2.0 fi = np.tan(alpha) In[0,0] = 1.0 In[1,1] = -1.0 #invert y In[2,2] = fi In[3,3] = 1.0 #to allow inversion return In def ray_to_normscreen(Ex,In,ray): """ ray_to_normscreen(Ex,In,ray) Ex - 4x4 extrinsic (modelview) matrix In - 4x4 projection matrix ray - 3x1 ray vector to project to the normalized (-1,1) x (-1,1) screen coordinates of a camera with a given orientation, position, and field of view described by Ex and In """ ray = ray.copy() rayp = np.zeros([4,1]) #projective ray rayp[0:3,0] = np.squeeze(ray) rayp[3,0] = 1.0 r2 = np.matmul(Ex,rayp) r3 = np.matmul(In,r2) if(r3[2,0]>0.0): sx = r3[0,0]/r3[2,0] sy = r3[1,0]/r3[2,0] else: sx = -2.0 sy = -2.0 #offscreen return [sx,sy] def vnorm(vect): vect = np.array(vect) vect = vect.copy() snrm = np.sqrt(np.sum(vect*vect)) return snrm def vnormalize(vect): vect = np.array(vect) vect = vect.copy() snrm = vnorm(vect) if(snrm>0.0): vect = vect/snrm else: vect = np.zeros(vect.shape) return vect def normscreen_to_ray(Ex,In,sx,sy): Ininv = np.linalg.inv(In) Exinv = np.linalg.inv(Ex) sp = np.zeros([4,1]) sp[0] = sx sp[1] = sy sp[2] = 1.0 sp[3] = 1.0 rp1 = np.matmul(Ininv,sp) rp2 = np.squeeze(rp1[0:3,0]) rp3 = vnormalize(rp2) rp4 = np.zeros([4,1]) rp4[0:3,0] = rp3[:] rp5 = np.matmul(Exinv,rp4) ray = np.squeeze(rp5[0:3,0]) return ray ### ### End vector math ### def _local_makeint(istr): try: i = int(istr) except: i = 0 return i def imfilefilter_preload(scenedir,flst): fnames = [] Nimgs = 0 angles = [] for f in flst: f2 = os.path.join(scenedir,f) fb = os.path.splitext(f)[0] fext = os.path.splitext(f)[1] if(fext.lower()==".png"): fsplit = fb.split("_") if(len(fsplit)==3): i1s = fsplit[1] i2s = fsplit[2] i1 = _local_makeint(i1s) i2 = _local_makeint(i2s) fnames.append(f2) angles.append([i1,0,i2]) Nimgs += 1 angles = np.float32(angles) return [Nimgs,angles,fnames] def load_all_images(scenedir, verbose=True): flst = os.listdir(scenedir) [Nimgs,angles,fnames] = imfilefilter_preload(scenedir,flst) imgs = None Nx = 0 Ny = 0 for I in range(0,Nimgs): if(verbose==True): print("Loading {} of {}".format(I,Nimgs)) img = load_img(fnames[I]) if(I==0): Nx = img.shape[0] Ny = img.shape[1] else: if(img.shape[0]!=Nx or img.shape[1]!=Ny): print("Error: Image {}, {} has size {},{}, not {},{}".format(I,fnames[I],img.shape[0],img.shape[1],Nx,Ny)) return None if(I==0): imgs = np.zeros((Nx,Ny,4,Nimgs),dtype=np.float32,order='F') imgs[:,:,:,I] = img[:,:,:] return [angles, imgs] def ray_to_normscreen(Ex,In,ray) """ ray_to_normscreen(Ex,In,ray) Is a given ray within the normalized screen coordinates of a camera described by Ex/In matrices? """ sx = -2.0 sy = -2.0 byesno = False [sx,sy] = ray_to_normscreen(Ex,In,ray) if(sx>=-1.0 and sx<=1.0 and sy>=-1.0 and sy<=1.0): byesno = True return [byesno, sx, sy] def check_cameraz(angles): #Debug script to check camera orientations #Camera z vectors should form a set covering the full spherical field of view if all the rotations are set up right Nimgs = angles.shape[0] zvs = np.zeros([3,Nimgs]) for I in range(0,Nimgs): R = rot_cameraxyz(angles[I,0]*pi/180,angles[I,1]*pi/180,angles[I,2]*pi/180) print("{} {} {}".format(angles[I,0],angles[I,1],angles[I,2])) print(R) zv = np.squeeze(R[0:3,2]) print(zv) zvs[:,I] = zv[:] # plt.figure() # plt.plot(angles[:,0],angles[:,2],'kd') fig = plt.figure() ax = fig.add_subplot(projection='3d') plt.plot(zvs[0,:],zvs[1,:],zvs[2,:],'kd') ax.set_xlim([-1,1]) ax.set_ylim([-1,1]) ax.set_zlim([-1,1]) plt.show() return def process_pixel(px,py,angles,imgs,fov,Nxout,Nyout): """ Eventually, figure out a way to vectorize this, or rewrite in C This takes a given output pixel px,py [0,Nxout) [0,Nyout) calculates a ray corresponding to the pixel for the lat-lon sky image and searches the available source images for which ones the ray lies within the field of view of. It pulls rgba from the source images and returns either an average or first available. """ Nx = imgs.shape[0] Ny = imgs.shape[1] Nimgs = imgs.shape[3] Nxc = np.arange(0,Nx) Nyc = np.arange(0,Ny) rgba = np.zeros(4) rgba_N = 0 rgba = np.float32([0,0,0,0]) ax = px/(Nxout)*2.0*pi ay = -2.0*(py/(Nyout-1)-0.5)*pi/2 ray = np.float32([cos(ax)*cos(ay),sin(ax)*cos(ay),sin(ay)]) # rgba = np.float32([ray[0],ray[1],ray[2],1]) #debug # rgba = 0.5*rgba+0.5 for I in range(0,Nimgs): R = rot_cameraxyz(angles[I,0]*pi/180,angles[I,1]*pi/180,angles[I,2]*pi/180) Ex = extrinsic(R,[0,0,0]) In = intrinsic_eqasp(fov) [byesno,sx,sy] = iswithin_img(ray,Ex,In) #debug #rgba = np.float32([ray[0]*0.5+0.5,ray[1]*0.5+0.5,0.5*ray[2]+0.5,1]) if(byesno): #interpolate images ra = imgs[:,:,0,I] ga = imgs[:,:,1,I] ba = imgs[:,:,2,I] aa = imgs[:,:,3,I] ppx = (sx+1.0)/2.0*(Nx-1.0) ppy = (-sy+1.0)/2.0*(Ny-1.0) ri = interp.interpn((Nxc,Nyc),ra,(ppx,ppy),bounds_error=False,fill_value=0.0) gi = interp.interpn((Nxc,Nyc),ga,(ppx,ppy),bounds_error=False,fill_value=0.0) bi = interp.interpn((Nxc,Nyc),ba,(ppx,ppy),bounds_error=False,fill_value=0.0) ai = interp.interpn((Nxc,Nyc),aa,(ppx,ppy),bounds_error=False,fill_value=0.0) rgba1 = np.squeeze(np.float32([ri,gi,bi,ai])) rgba = rgba + rgba1 rgba_N += 1 #debug #rgba = np.float32([0.5*sx+0.5,0.5*sy+0.5,0,1]) #rgba = np.float32([angles[I,0]/360,0,0,1]) break if(rgba_N>0): rgba = rgba/rgba_N return rgba def proc_pixelrange(angles,sm_imgs,imgs_shape,fov,sm_imgout,imgout_shape,threadnum,nthreads): """ Thread function for processing a range of pixels. angles - the list of camera angles sm_imgs - the shared memory buffer for the source images imgs_shape - the shape of the source images fov - the field of view angle for the camera source images [rad] sm_imgout - the shared memory buffer for the output images imgout_shape - the shape of the output image buffer threadnum - which thread is this? nthreads - total number of threads """ Nxout = imgout_shape[0] Nyout = imgout_shape[1] #Thread indexing #Is - stride index - how many elements should each thread process? #I0,I1 - elements processed by this thread #N - total number of pixels N = Nxout*Nyout Is = int(np.floor(N/nthreads))+1 I0 = threadnum*Is I1 = (threadnum+1)*Is if(I0>N): I0 = N if(I1>N): I1 = N #Unpack shared memory arrays s_imgs = np.ndarray(imgs_shape,dtype=np.float32,buffer=sm_imgs.buf) s_imgout = np.ndarray(imgout_shape,dtype=np.float32,buffer=sm_imgout.buf) for I in range(I0,I1): py = int(np.floor(I/Nxout)) px = I%Nxout rgba = process_pixel(px,py,angles,s_imgs,fov,Nxout,Nyout) s_imgout[px,py,:] = rgba.flatten()[:] #save rgba to shared-memory output array return def threaded_proc(angles,imgs,fov,imgout): """ Thread manager for parallel processing. The actual threading library doesn't work for this purpose because of GIL multiprocessing does allow parallel execution, so "threads" are "multiprocessing" Processes Responsibility for painting in pixels is split among a number of threads. While this is faster than single threaded, it is still pretty slow by compiled code or GPU code standards, and I'll probably write a compiled submodule for this in the future. For portability's sake, though, here is the native python version of this. """ nthreads = os.cpu_count()-2 if(nthreads<=0): nthreads=1 threads = [] #create shared memory arrays sm_imgs = shared_memory.SharedMemory(create=True, size=imgs.nbytes) s_imgs = np.ndarray(imgs.shape,dtype=np.float32,buffer=sm_imgs.buf) s_imgs[:] = imgs[:] sm_imgout = shared_memory.SharedMemory(create=True, size=imgout.nbytes) s_imgout = np.ndarray(imgout.shape,dtype=np.float32,buffer=sm_imgout.buf) for I in range(0,nthreads): t = mp.Process(target=proc_pixelrange,args=(angles,sm_imgs,s_imgs.shape,fov,sm_imgout,s_imgout.shape,I,nthreads)) threads.append(t) for I in range(0,nthreads): threads[I].start() for I in range(0,nthreads): threads[I].join() print("{} %% complete".format((I+1)/nthreads*100)) imgout[:] = s_imgout[:] sm_imgs.close() sm_imgs.unlink() sm_imgout.close() sm_imgout.unlink() return imgout def main(): [angles,imgs] = load_all_images("./scenerenders") fnout = "./scenehdr2.png" Nxout = 200 #Dimensions of the output scene image Nyout = int(np.floor(Nxout/2.0)) fov = 50.0*pi/180.0 #field of view angle of the renders in [radians] imgout = np.zeros((Nxout,Nyout,4),dtype=np.float32) # single threaded implementation # for I in range(0,Nxout): # print("processing {:1.2f} %% complete.".format(I/Nxout*100.0)) # for J in range(0,Nyout): # rgba = process_pixel(I,J,angles,imgs,fov,Nxout,Nyout) # imgout[I,J,:] = rgba[:] imgout = threaded_proc(angles,imgs,fov,imgout) save_img(fnout,imgout) return if(__name__=="__main__"): main() #test_eulermatrices() exit(0)