def main(): parser = argparse.ArgumentParser() parser.add_argument('--cfg_file', type=str,default = './configs/shapenet/shapenet_1c_02691156.yaml') parser.add_argument('--checkpoint', type=str,default = './data/models/02691156-R.pth.tar') parser.add_argument('--init_obj', type=str,default = './data/sphere_642.obj') parser.add_argument('--input_path', type=str,default = './data/demo/') args = parser.parse_args() # load cfg f = open(args.cfg_file) cfg = yaml.load(f) f.close() # make model model_R = models.Reconstructor(args.init_obj,cfg).to(device) checkpoint = torch.load(args.checkpoint) model_R.load_state_dict(checkpoint['model_R']) model_R.eval() img_files = os.listdir(args.input_path) for i in range(len(img_files)): print(str(i) + '/' + str(len(img_files))) temp_img_path = os.path.join(args.input_path , img_files[i]) filename, file_extension = os.path.splitext(img_files[i]) if not (file_extension == '.png' or file_extension == '.jpg'): continue temp_img = misc.imread(temp_img_path) temp_img = temp_img.transpose(2,0,1) img = torch.FloatTensor(temp_img.astype('float32') / 255.).unsqueeze(0).to(device) vertices, faces,_ = model_R.reconstruct(img) srf.save_obj(os.path.join(args.input_path,filename+'.obj'),vertices[0],faces[0])
def test(): end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() losses1 = AverageMeter() iou_all = [] for class_id, class_name in dataset_val.class_ids_pair: directory_mesh_cls = os.path.join(directory_mesh, class_id) os.makedirs(directory_mesh_cls, exist_ok=True) iou = 0 for i, (im, vx) in enumerate( dataset_val.get_all_batches_for_evaluation( args.batch_size, class_id)): images = torch.autograd.Variable(im).cuda() voxels = vx.numpy() batch_iou, vertices, faces = model(images, voxels=voxels, task='test') iou += batch_iou.sum() batch_time.update(time.time() - end) end = time.time() # save demo images for k in range(vertices.size(0)): obj_id = (i * args.batch_size + k) if obj_id % args.save_freq == 0: mesh_path = os.path.join(directory_mesh_cls, '%06d.obj' % obj_id) input_path = os.path.join(directory_mesh_cls, '%06d.png' % obj_id) srf.save_obj(mesh_path, vertices[k], faces[k]) imageio.imsave(input_path, img_cvt(images[k])) # print loss if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time.val:.3f}\t' 'IoU {2:.3f}\t'.format( i, ((dataset_val.num_data[class_id] * 24) // args.batch_size), batch_iou.mean(), batch_time=batch_time)) iou_cls = iou / 24. / dataset_val.num_data[class_id] * 100 iou_all.append(iou_cls) print('=================================') print('Mean IoU: %.3f for class %s' % (iou_cls, class_name)) print('\n') print('=================================') print('Mean IoU: %.3f for all classes' % (sum(iou_all) / len(iou_all)))
def save_obj(self, filename_obj, save_texture=False, texture_res_out=16): if self.batch_size != 1: raise ValueError('Could not save when batch size >= 1') if save_texture: srf.save_obj(filename_obj, self.vertices[0], self.faces[0], textures=self.textures[0], texture_res=texture_res_out, texture_type=self.texture_type) else: srf.save_obj(filename_obj, self.vertices[0], self.faces[0], textures=None)
for each_key in category_to_IoU: save_gen_model_path = os.path.join(save_dir, each_key + '_generated_model.obj') save_ori_model_path = os.path.join(save_dir, each_key + '_original_model.obj') file_list = sorted(os.listdir(os.path.join(root_dir, each_key))) if '000000__broken__67ada28ebc79cc75a056f196c127ed77' in file_list: file_list.remove('000000__broken__67ada28ebc79cc75a056f196c127ed77') if '000000__broken__b65b590a565fa2547e1c85c5c15da7fb' in file_list: file_list.remove('000000__broken__b65b590a565fa2547e1c85c5c15da7fb') image_path = os.path.join(root_dir, each_key, file_list[category_to_argmin[each_key]], 'rendered_0.png') original_model_path = os.path.join(root_dir, each_key, file_list[category_to_argmin[each_key]], 'model.obj') image = PIL.Image.open(image_path) image = image.resize((64, 64)) image = np.asanyarray(image)[None, :, :, :] image = image.transpose((0, 3, 1, 2)) image = np.ascontiguousarray(image) image = torch.from_numpy(image.astype('float32') / 255.) image = torch.autograd.Variable(image).cuda() vertice_generate, face_generate = model.reconstruct(image) srf.save_obj(save_gen_model_path, vertice_generate[0], face_generate[0]) shutil.copy(original_model_path, save_ori_model_path) print(each_key, category_to_num[each_key], category_to_max[each_key], category_to_min[each_key], category_to_mean[each_key])
def train(): writer = SummaryWriter() print("Made summary writer") end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() for i in range(start_iter, args.num_iterations + 1): # adjust learning rate and sigma_val (decay after 150k iter) lr = adjust_learning_rate([optimizer], args.learning_rate, i, method=args.lr_type) model.set_sigma(adjust_sigma(args.sigma_val, i)) # TRAIN # load images from multi-view images_a, images_b, viewpoints_a, viewpoints_b = dataset_train.get_random_batch( args.batch_size) images_a = images_a.cuda() images_b = images_b.cuda() viewpoints_a = viewpoints_a.cuda() viewpoints_b = viewpoints_b.cuda() # soft render images render_images, laplacian_loss, flatten_loss = model( [images_a, images_b], [viewpoints_a, viewpoints_b], task='train') laplacian_loss = laplacian_loss.mean() flatten_loss = flatten_loss.mean() # compute loss loss = multiview_iou_loss(render_images, images_a, images_b) + \ args.lambda_laplacian * laplacian_loss + \ args.lambda_flatten * flatten_loss # VAL # load images from multi-view val_images_a, val_images_b, val_viewpoints_a, val_viewpoints_b = dataset_val.get_random_batch( args.batch_size) val_images_a = val_images_a.cuda() val_images_b = val_images_b.cuda() val_viewpoints_a = val_viewpoints_a.cuda() val_viewpoints_b = val_viewpoints_b.cuda() # soft render images val_render_images, val_laplacian_loss, val_flatten_loss = model( [val_images_a, val_images_b], [val_viewpoints_a, val_viewpoints_b], task='train') val_laplacian_loss = val_laplacian_loss.mean() val_flatten_loss = val_flatten_loss.mean() # compute loss val_loss = multiview_iou_loss(val_render_images, val_images_a, val_images_b) + \ args.lambda_laplacian * val_laplacian_loss + \ args.lambda_flatten * val_flatten_loss # RECORD LOSSES writer.add_scalar('Loss/train', loss.data.item(), i) writer.add_scalar('Loss/val', val_loss.data.item(), i) losses.update(loss.data.item(), images_a.size(0)) # compute gradient and optimize optimizer.zero_grad() loss.backward() optimizer.step() batch_time.update(time.time() - end) end = time.time() # save checkpoint if i % args.save_freq == 0: model_path = os.path.join(directory_output, 'checkpoint_%07d.pth.tar' % i) torch.save( { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), }, model_path) # save demo images if i % args.demo_freq == 0: demo_image = images_a[0:1] demo_path = os.path.join(directory_output, 'demo_%07d.obj' % i) demo_v, demo_f = model.reconstruct(demo_image) srf.save_obj(demo_path, demo_v[0], demo_f[0]) fake_img = img_cvt(render_images[0][0]) real_img = img_cvt(images_a[0]) imageio.imsave(os.path.join(image_output, '%07d_fake.png' % i), fake_img) imageio.imsave(os.path.join(image_output, '%07d_input.png' % i), real_img) # writer.add_image('Fake image', fake_img, i) # writer.add_image('Real image', real_img, i) # print if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time.val:.3f}\t' 'Loss {loss.val:.3f}\t' 'lr {lr:.6f}\t' 'sv {sv:.6f}\t'.format( i, args.num_iterations, batch_time=batch_time, loss=losses, lr=lr, sv=model.renderer.rasterizer.sigma_val)) writer.close()
def train(): end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() iou_losses = [] for i in range(start_iter, args.num_iterations + 1): # adjust learning rate and sigma_val (decay after 150k iter) lr = adjust_learning_rate([optimizer], args.learning_rate, i, method=args.lr_type) model.set_sigma(adjust_sigma(args.sigma_val, i)) # load images from multi-view images_a, images_b, viewpoints_a, viewpoints_b = dataset_train.get_random_batch( args.batch_size) images_a = images_a.cuda() images_b = images_b.cuda() viewpoints_a = viewpoints_a.cuda() viewpoints_b = viewpoints_b.cuda() # soft render images render_images, laplacian_loss, flatten_loss = model( [images_a, images_b], [viewpoints_a, viewpoints_b], task='train') laplacian_loss = laplacian_loss.mean() flatten_loss = flatten_loss.mean() iou_loss = multiview_iou_loss(render_images, images_a, images_b) iou_losses.append(iou_loss.data.item()) # compute loss loss = iou_loss + \ args.lambda_laplacian * laplacian_loss + \ args.lambda_flatten * flatten_loss losses.update(loss.data.item(), images_a.size(0)) # compute gradient and optimize optimizer.zero_grad() loss.backward() optimizer.step() batch_time.update(time.time() - end) end = time.time() # save checkpoint if i % args.save_freq == 0: model_path = os.path.join(directory_output, 'checkpoint_%07d.pth.tar' % i) torch.save( { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), }, model_path) # save demo images if i % args.demo_freq == 0: demo_image = images_a[0:1] demo_path = os.path.join(directory_output, 'demo_%07d.obj' % i) demo_v, demo_f = model.reconstruct(demo_image) srf.save_obj(demo_path, demo_v[0], demo_f[0]) imageio.imsave(os.path.join(image_output, '%07d_fake.png' % i), img_cvt(render_images[0][0])) imageio.imsave(os.path.join(image_output, '%07d_input.png' % i), img_cvt(images_a[0])) # print if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time.val:.3f}\t' 'Loss {loss.val:.3f}\t' 'lr {lr:.6f}\t' 'sv {sv:.6f}\t'.format( i, args.num_iterations, batch_time=batch_time, loss=losses, lr=lr, sv=model.renderer.rasterizer.sigma_val)) print(iou_losses) with open('loss.txt', 'w') as f: f.write(str(iou_losses))
def train(): end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() env_name = args.experiment_id img_a_logger = VisdomLogger('images', env=env_name, port=8097, opts=dict(title='img_a')) img_b_logger = VisdomLogger('images', env=env_name, port=8097, opts=dict(title='img_b')) rnd_logger = [VisdomLogger('images', env=env_name, port=8097, opts=dict(title='rnd_rgb_{}'.format(_i))) for _i in range(4)] for i in range(start_iter, args.num_iterations + 1): # adjust learning rate and sigma_val (decay after 150k iter) lr = adjust_learning_rate([optimizer], args.learning_rate, i, method=args.lr_type) model.set_sigma(adjust_sigma(args.sigma_val, i)) # load images from multi-view images_a, images_b, viewpoints_a, viewpoints_b = dataset_train.get_random_batch(args.batch_size) images_a = images_a.cuda() images_b = images_b.cuda() viewpoints_a = viewpoints_a.cuda() viewpoints_b = viewpoints_b.cuda() # soft render images render_images, laplacian_loss, flatten_loss = model([images_a, images_b], [viewpoints_a, viewpoints_b], task='train') img_a_logger.log(F.interpolate(images_a.detach().clone(), scale_factor=4).squeeze()[:,:3]) img_b_logger.log(F.interpolate(images_b.detach().clone(), scale_factor=4).squeeze()[:,:3]) for _i in range(4): rnd_logger[_i].log(F.interpolate(render_images[_i].detach().clone(), scale_factor=4).squeeze()[:,:3]) import ipdb; ipdb.set_trace() laplacian_loss = laplacian_loss.mean() flatten_loss = flatten_loss.mean() # compute loss loss = multiview_iou_loss(render_images, images_a, images_b) + \ args.lambda_laplacian * laplacian_loss + \ args.lambda_flatten * flatten_loss losses.update(loss.data.item(), images_a.size(0)) # compute gradient and optimize optimizer.zero_grad() loss.backward() optimizer.step() batch_time.update(time.time() - end) end = time.time() # save checkpoint if i % args.save_freq == 0: model_path = os.path.join(directory_output, 'checkpoint_%07d.pth.tar'%i) torch.save({ 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), }, model_path) # save demo images if i % args.demo_freq == 0: demo_image = images_a[0:1] demo_path = os.path.join(directory_output, 'demo_%07d.obj'%i) demo_v, demo_f = model.reconstruct(demo_image) srf.save_obj(demo_path, demo_v[0], demo_f[0]) imageio.imsave(os.path.join(image_output, '%07d_fake.png' % i), img_cvt(render_images[0][0])) imageio.imsave(os.path.join(image_output, '%07d_input.png' % i), img_cvt(images_a[0])) # print if i % args.print_freq == 0: print('Iter: [{0}/{1}]\t' 'Time {batch_time.val:.3f}\t' 'Loss {loss.val:.3f}\t' 'lr {lr:.6f}\t' 'sv {sv:.6f}\t'.format(i, args.num_iterations, batch_time=batch_time, loss=losses, lr=lr, sv=model.renderer.rasterizer.sigma_val))
end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() losses1 = AverageMeter() iou_all = [] images = [] img_list = sorted(os.listdir(IMG_PATH)) for img_name in img_list: img = PIL.Image.open(os.path.join(IMG_PATH, img_name)) img = np.asanyarray(img) images.append(img) images = np.array(images) images = images.transpose((0, 3, 1, 2)) images = np.ascontiguousarray(images) images = torch.from_numpy(images.astype('float32') / 255.) images = torch.autograd.Variable(images).cuda() vertices, faces = model.reconstruct(images) for k in range(len(img_list)): print(k) mesh_path = os.path.join(directory_output, img_list[k][:-4] + ".obj") srf.save_obj(mesh_path, vertices[k], faces[k])