Spaces:
Runtime error
Runtime error
File size: 4,234 Bytes
c7f097c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
import time
import json
import numpy as np
import torch
from torch.utils.data import DataLoader
from lib.options import BaseOptions
from lib.mesh_util import *
from lib.sample_util import *
from lib.train_util import *
from lib.model import *
from PIL import Image
import torchvision.transforms as transforms
import glob
import tqdm
# get options
opt = BaseOptions().parse()
class Evaluator:
def __init__(self, opt, projection_mode='orthogonal'):
self.opt = opt
self.load_size = self.opt.loadSize
self.to_tensor = transforms.Compose([
transforms.Resize(self.load_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# set cuda
cuda = torch.device('cuda:%d' % opt.gpu_id) if torch.cuda.is_available() else torch.device('cpu')
# create net
netG = HGPIFuNet(opt, projection_mode).to(device=cuda)
print('Using Network: ', netG.name)
if opt.load_netG_checkpoint_path:
netG.load_state_dict(torch.load(opt.load_netG_checkpoint_path, map_location=cuda))
if opt.load_netC_checkpoint_path is not None:
print('loading for net C ...', opt.load_netC_checkpoint_path)
netC = ResBlkPIFuNet(opt).to(device=cuda)
netC.load_state_dict(torch.load(opt.load_netC_checkpoint_path, map_location=cuda))
else:
netC = None
os.makedirs(opt.results_path, exist_ok=True)
os.makedirs('%s/%s' % (opt.results_path, opt.name), exist_ok=True)
opt_log = os.path.join(opt.results_path, opt.name, 'opt.txt')
with open(opt_log, 'w') as outfile:
outfile.write(json.dumps(vars(opt), indent=2))
self.cuda = cuda
self.netG = netG
self.netC = netC
def load_image(self, image_path, mask_path):
# Name
img_name = os.path.splitext(os.path.basename(image_path))[0]
# Calib
B_MIN = np.array([-1, -1, -1])
B_MAX = np.array([1, 1, 1])
projection_matrix = np.identity(4)
projection_matrix[1, 1] = -1
calib = torch.Tensor(projection_matrix).float()
# Mask
mask = Image.open(mask_path).convert('L')
mask = transforms.Resize(self.load_size)(mask)
mask = transforms.ToTensor()(mask).float()
# image
image = Image.open(image_path).convert('RGB')
image = self.to_tensor(image)
image = mask.expand_as(image) * image
return {
'name': img_name,
'img': image.unsqueeze(0),
'calib': calib.unsqueeze(0),
'mask': mask.unsqueeze(0),
'b_min': B_MIN,
'b_max': B_MAX,
}
def eval(self, data, use_octree=False):
'''
Evaluate a data point
:param data: a dict containing at least ['name'], ['image'], ['calib'], ['b_min'] and ['b_max'] tensors.
:return:
'''
opt = self.opt
with torch.no_grad():
self.netG.eval()
if self.netC:
self.netC.eval()
save_path = '%s/%s/result_%s.obj' % (opt.results_path, opt.name, data['name'])
if self.netC:
gen_mesh_color(opt, self.netG, self.netC, self.cuda, data, save_path, use_octree=use_octree)
else:
gen_mesh(opt, self.netG, self.cuda, data, save_path, use_octree=use_octree)
if __name__ == '__main__':
evaluator = Evaluator(opt)
test_images = glob.glob(os.path.join(opt.test_folder_path, '*'))
test_images = [f for f in test_images if ('png' in f or 'jpg' in f) and (not 'mask' in f)]
test_masks = [f[:-4]+'_mask.png' for f in test_images]
print("num; ", len(test_masks))
for image_path, mask_path in tqdm.tqdm(zip(test_images, test_masks)):
try:
print(image_path, mask_path)
data = evaluator.load_image(image_path, mask_path)
evaluator.eval(data, True)
except Exception as e:
print("error:", e.args)
|