data_prep.py

import os
import sys

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
sys.path.append(BASE_DIR + '/scripts/')
import urllib.request 
from multiprocessing import Pool
import binvox_rw
import scipy.io as sio
import numpy as np
from tqdm import tqdm
from glob import glob
import random
import shutil
from PIL import Image
import argparse
from scipy import ndimage
from subprocess import call
import torch
import utils
from time import time

parser = argparse.ArgumentParser(description='Dataset prep for image to 3D object super resolution')
parser.add_argument('-no','--num_objects', default=500, help='number of objects to be converted', type = int)
args = parser.parse_args()


#labels for the union of the core shapenet classes and the ikea dataset classes
labels = {'04379243':'table','03211117':'monitor','04401088':'cellphone','04530566': 'watercraft',  '03001627' : 'chair','03636649' : 'lamp',  '03691459': 'speaker' ,  '02828884':'bench',
'02691156': 'plane', '02808440': 'bathtub',  '02871439': 'bookcase',
'02773838': 'bag', '02801938': 'basket', '02828884' : 'bench','02880940': 'bowl' ,
'02924116': 'bus', '02933112': 'cabinet', '02942699': 'camera', '02958343': 'car', '03207941': 'dishwasher',
'03337140': 'file', '03624134': 'knife', '03642806': 'laptop', '03710193': 'mailbox',
'03761084': 'microwave', '03928116': 'piano', '03938244':'pillow', '03948459': 'pistol', '04004475': 'printer',
'04099429': 'rocket', '04256520': 'sofa', '04554684': 'washer', '04090263': 'rifle'}

objects = ['bench','cabinet','car','cellphone','chair','lamp','monitor','plane','rifle','sofa','speaker','table','watercraft']
wanted_classes=[]
for l in labels:
	if labels[l] in objects:
		wanted_classes.append(l)

debug_mode = False # change to make all of the called scripts print their errors and warnings
if debug_mode:
	io_redirect = ''
else:
	io_redirect = ' > /dev/null 2>&1'


# make data directories
if not os.path.exists('data/objects/'):
	os.makedirs('data/objects/')



# download .obj obect files
def download():
	with open('scripts/binvox_file_locations.txt','rb') as f: # location of all the binvoxes for shapenet's core classes
		content = f.readlines()

	# make data sub-directories for each class
	for s in wanted_classes:
		obj = 'data/objects/' + labels[s]+'/'
		if not os.path.exists(obj):
			os.makedirs(obj)


	# search object for correct object classes
	binvox_urls = []
	obj_urls = []
	for file in content:
		file = str(file)
		current_class = file.split('/')
		if current_class[1] in wanted_classes:
			if '_' in current_class[3]: continue
			if 'presolid' in current_class[3]: continue
			obj_urls.append(['http://shapenet.cs.stanford.edu/shapenet/obj-zip/ShapeNetCore.v1/'+file.split('/')[1]+'/'+file.split('/')[2]+'/model.obj', 'data/objects/'+labels[current_class[1]]+ '/'+ current_class[2]+'.obj'])

	# get randomized sample from each object class of correct size
	random.shuffle(obj_urls)
	final_urls = []
	dictionary = {}
	for o in obj_urls:
		obj_class = o[1].split('/')[-2]
		if obj_class in dictionary:
			dictionary[obj_class] += 1
			if dictionary[obj_class]> args.num_objects:
				continue
		else:
			dictionary[obj_class] = 1
		final_urls.append(o)

	# parallel downloading of object .obj files
	pool = Pool(processes=16)
	pbar = tqdm(pool.imap_unordered(down, final_urls), total=len(final_urls))
	pbar.set_description(f"Downloading Meshes")
	for _ in pbar:
		pass
	

# this take object files and makes then a more managable size
# this is only done for training the latent loss
# it makes it far quicker to load the object during training
def manage_objects():
	commands = []
	for s in wanted_classes:
		# final all downloaded objects from the class
		objs = glob('data/objects/' + labels[s]+'/*.obj')
		location_meshinfo = 'data/mesh_info/' + labels[s]+'/'
		location_obj = 'data/managable_objects/' + labels[s]+'/'
		if not os.path.exists(location_meshinfo):
			os.makedirs(location_meshinfo)
		if not os.path.exists(location_obj):
			os.makedirs(location_obj)
		l = 0

		for o in objs:
			name = o.split('/')[-1][:-4]
			file_name_mesh =  location_meshinfo+ name
			file_name_new_obj = location_obj + name + '.obj'
			cmd = 'blender scripts/manage.blend -b -P scripts/blender_convert.py -- %s %s %s' %( o, file_name_mesh,file_name_new_obj )
			commands.append(cmd)


	random.shuffle(commands)
	pool = Pool(processes=16)
	pbar = tqdm(pool.imap_unordered(call, commands), total=len(commands))
	pbar.set_description(f"Downscaling meshes")
	for _ in pbar:
		pass

	message = 'The blender commands failed. Please check why using the following command: '
	assert len(glob('data/managable_objects/' + labels[s]+'/*.obj')) > 0, message + cmd


# converts obj files to binvox, is an intermediary for voxel computation
def binvox():
	commands =[]
	for s in wanted_classes:
		dirs = glob('data/managable_objects/' + labels[s]+'/*.obj')
		
		count = 0
		for d in (dirs):
			command = 'scripts/binvox ' + d  + ' -d ' + str(32)+ ' -pb -cb -c -e'   # this executable can be found at http://www.patrickmin.com/binvox/ ,
			# -d x idicates resoltuion will be x by x by x , -pb is to stop the visualization, the rest of the commnads are to help make the object water tight
			commands.append(command)
		
	
	random.shuffle(commands)
	pool = Pool(processes=16)
	pbar = tqdm(pool.imap_unordered(call, commands), total=len(commands))
	pbar.set_description(f"converting meshes to small binvoxes")
	for _ in pbar:
		pass

	message = 'The binvox executable failed. Please check its permissions, and that it can run properly from the commandline using the following command: '
	assert len(glob('data/managable_objects/' + labels[s]+'/*.binvox')) > 0, message + command
		


# converts binvox files to voxel files
def convert_bin():
	models = []
	for s in wanted_classes:
		directory = 'data/voxels/'+labels[s] +'/'
		# find all binvoxes
		models  += glob('data/managable_objects/'+labels[s]+'/*.binvox')
		if not os.path.exists(directory):
			os.makedirs(directory)

	cur_class = ''
	random.shuffle(models)
	pbar = tqdm(models)
	pbar.set_description(f"Converting binvoxes to voxels")
	for m in pbar:
		with open(m, 'rb') as f:
			try:
				model = binvox_rw.read_as_3d_array(f).data
			except ValueError:
				continue
		directory = 'data/voxels/' + m.split('/')[-2] + '/'
			
		# remove internals from models
		# I think this makes it easier to learn
		positions = np.where(model != 0 )
		new_mod = np.zeros(model.shape)
		for i, j, k in zip(*positions):
			# identifies if current voxel has an exposed face
			if np.sum(model[i-1:i+2, j-1:j+2, k-1:k+2]) < 27:
				new_mod[i,j,k] = 1
		# save as np array

		sio.savemat(directory +m.split('/')[-1][:-7], {'model': new_mod.astype(np.uint8)})





# these are two simple functions for parallel processing
# down() downloads , and call() calls functions
def down(url):
	urllib.request.urlretrieve(url[0], url[1])
def call(command):
	os.system('%s %s' % (command, io_redirect))




# calculates the surface of an object
def surface_computation(data):
	dim = data.shape[0]

	a,b,c = np.where(data == 1)
	large = int(dim *1.5)
	big_list = [[[[-1,large]for j in range(dim)] for i in range(dim)] for k in range(3)]
	# over the whole object extract for each face the first and last occurance of a voxel at each pixel
	# we take highest for convinience
	for i,j,k in zip(a,b,c):
		big_list[0][i][j][0] = (max(k,big_list[0][i][j][0]))
		big_list[0][i][j][1] = (min(k,big_list[0][i][j][1]))
		big_list[1][i][k][0] = (max(j,big_list[1][i][k][0]))
		big_list[1][i][k][1] = (min(j,big_list[1][i][k][1]))
		big_list[2][j][k][0] = (max(i,big_list[2][j][k][0]))
		big_list[2][j][k][1] = (min(i,big_list[2][j][k][1]))
	faces = np.zeros((6,dim,dim)) # will hold odms
	for i in range(dim):
		for j in range(dim):
			faces[0,i,j] =   dim -1 - big_list[0][i][j][0]         if    big_list[0][i][j][0]   > -1 else dim
			# we subtract from the (dimension -1) as we computed the last occurance, instead of the first for half of the faces
			faces[1,i,j] =   big_list[0][i][j][1]        		   if    big_list[0][i][j][1]   < large else dim
			faces[2,i,j] =   dim -1 - big_list[1][i][j][0]         if    big_list[1][i][j][0]   > -1 else dim
			faces[3,i,j] =   big_list[1][i][j][1]        		   if    big_list[1][i][j][1]   < large else dim
			faces[4,i,j] =   dim -1 - big_list[2][i][j][0]         if    big_list[2][i][j][0]   > -1 else dim
			faces[5,i,j] =   big_list[2][i][j][1]         		   if    big_list[2][i][j][1]   < large else dim
	return faces






def calc_surface():
	dims = 128
	# first render as voxel array at high resolution
	models = []
	for s in wanted_classes:
		models += glob('data/objects/' + labels[s]+'/*.obj')

	commands =[]
	random.shuffle(models)
	for m in models:
		command = 'scripts/binvox ' + m  + ' -d ' + str(dims)+ ' -pb -cb -c -e'   # this executable can be found at http://www.patrickmin.com/binvox/ ,
		# -d x idicates resoltuion will be x by x by x , -pb is to stop the visualization, the rest of the commnads are to help make the object water tight
		commands.append(command)

	pool = Pool(processes=16)
	pbar = tqdm(pool.imap_unordered(call, commands), total=len(commands))
	pbar.set_description(f"Making large binvoxes from meshes")
	for _ in pbar:
		pass

	message = 'The binvox executable failed. Please check its permissions, and that it can run properly from the commandline using the following command: '
	assert len(glob('data/objects/' + labels[s]+'/*.binvox')) > 0, message + command
		

	

	models = []
	for s in wanted_classes:

		models  += glob('data/objects/'+labels[s]+'/*.binvox')
		location = 'data/surfaces/'+labels[s] +'/'
		if not os.path.exists(location):
			os.makedirs(location)
	random.shuffle(models)
	pbar = tqdm(models)
	pbar.set_description(f"Extracting point cloud from converted objects")
	for m in pbar:
		location = 'data/surfaces/'+m.split('/')[-2] +'/'
		with open(m, 'rb') as f:
			try:
				model = binvox_rw.read_as_3d_array(f).data
			except ValueError:
				continue
		faces = surface_computation(model)
		# calculate low resolution version, to make watertight
		high, low = dims, 32
		down = high // low
		a,b,c = np.where(model==1)
		low_model = np.zeros((low,low,low))
		for x,y,z in zip(a,b,c):
				low_model[ x//down, y//down, z//down] =1
		# fill internals
		low_model[ndimage.binary_fill_holes(low_model)] = 1

		# obtain surface projections
		corrected = np.zeros((high,high,high))
		for i in range(low):
			for j in range(low):
				for k in range(low):
					corrected[i*down: (i+1)*down, j*down:(j+1)*down, k*down:(k+1)*down] = low_model[i,j,k]
		# carve away from low res model
		for i in range(high):
			for j in range(high):
				if faces[0,i,j] <high:
					corrected[i,j,int((high - faces[0,i,j])):high]=0
				else:
					corrected[i,j,:] =0

				if faces[1,i,j] <high:
					corrected[i,j,0:int(faces[1,i,j])]=0
				else:
					corrected[i,j,:] =0

				if faces[2,i,j] <high:
					corrected[i,int((high - faces[2,i,j])):high, j] =0
				else:
					corrected[i,:,j] =0

				if faces[3,i,j] <high:
					corrected[i,0:int(faces[3,i,j]), j] =0
				else:
					corrected[i,:,j] =0

				if faces[4,i,j] <high:
					corrected[int((high - faces[4,i,j])):high,i,j] =0
				else:
					corrected[:,i,j] =0

				if faces[5,i,j] <high:
					corrected[0:int(faces[5,i,j]),i,j] =0
				else:
					corrected[:,i,j] =0

		corrected[ndimage.binary_fill_holes(corrected)] = 1
		positions = np.where(corrected != 0 )
		new_mod = np.zeros(corrected.shape)
		points = []
		# get only surface objects
		for i, j, k in zip(*positions):
			# identifies if current voxel has an exposed face
			if np.sum(corrected[i-1:i+2, j-1:j+2, k-1:k+2]) < 27:
				points.append([i,j,k])
		voxel_points = np.array(points).astype(float)

		obj = m[:-7] + '.obj'
		try:
			obj = utils.ObjLoader(obj)
		except ValueError:
			continue

		voxel_points = np.array(voxel_points)
		mesh_points = np.array(obj.vertices)

		# make computed surface be same size as origional obejct
		xx = np.amax(mesh_points[:,0]) - np.amin(mesh_points[:,0])
		xx_v = np.amax(voxel_points[:,0]) - np.amin(voxel_points[:,0])
		x_diff = xx/xx_v

		yy = np.amax(mesh_points[:,1]) - np.amin(mesh_points[:,1])
		yy_v = np.amax(voxel_points[:,1]) - np.amin(voxel_points[:,1])
		y_diff = yy/yy_v

		zz = np.amax(mesh_points[:,2]) - np.amin(mesh_points[:,2])
		zz_v = np.amax(voxel_points[:,2]) - np.amin(voxel_points[:,2])
		z_diff = zz/zz_v
		voxel_points*=[x_diff, y_diff, z_diff]

		xx = np.amax(mesh_points[:,0])
		xx_v = np.amax(voxel_points[:,0])
		x_diff = xx-xx_v

		yy = np.amax(mesh_points[:,1])
		yy_v = np.amax(voxel_points[:,1])
		y_diff = yy-yy_v

		zz = np.amax(mesh_points[:,2])
		zz_v = np.amax(voxel_points[:,2])
		z_diff = zz-zz_v
	

		voxel_points+=[x_diff, y_diff, z_diff]
		
		while voxel_points.shape[0] < 10000: 
			voxel_points = np.concatenate((voxel_points, voxel_points))
		sio.savemat(location + m.split('/')[-1][:-7] , {'points': voxel_points})

def download_images():
	print('downloading shapenet images')
	command = 'wget http://cvgl.stanford.edu/data2/ShapeNetRendering.tgz -O data/images.tgz'
	os.system(command)
	
	print('extracting shapenet images')
	command = 'tar -xzf data/images.tgz --directory data/'
	os.system(command)
	
	os.rename('data/ShapeNetRendering','data/images' )

	print('Splitting dataset in training, valudation, and tests.')
	folders = glob('data/images/*/')
	for folder in tqdm(folders): 

		im_fol = glob( folder + '*')

		if not os.path.exists(folder + 'train'):
			os.mkdir(folder + 'train')
			os.mkdir(folder + 'valid')
			os.mkdir(folder + 'test')

		first = (7*len(im_fol))//10
		second  = (8*len(im_fol))//10
		
		train = im_fol[:first]
		valid = im_fol[first:second]
		test = im_fol[second:]

		for ex in train: 
			source = ex 
			dest = folder + 'train/' + ex.split('/')[-1]
			
			shutil.move(source, dest)

		for  ex in valid: 
			source = ex 
			dest = folder + 'valid/' + ex.split('/')[-1]
			
			shutil.move(source, dest)

		for  ex in test: 
			source = ex 
			dest = folder + 'test/' + ex.split('/')[-1]
			
			shutil.move(source, dest)
	for folder in folders:
		class_num = folder.split('/')[-2]
		class_obj = labels[class_num]
		os.rename(folder, folder[:-len(class_num)-1] + class_obj)







download()
manage_objects()
binvox()
convert_bin()
calc_surface()
download_images()
print ('finished eratin')