def run_viz_suite(name,
data,
dest_task=tasks.depth_zbuffer,
graph_file=None,
model_file=None,
logger=None,
old=False,
multitask=False,
percep_mode=None):
if graph_file is not None:
graph = TaskGraph(tasks=[tasks.rgb, dest_task], pretrained=False)
graph.load_weights(graph_file)
model = graph.edge(tasks.rgb, dest_task).load_model()
elif old:
model = DataParallelModel.load(UNetOld().cuda(), model_file)
elif multitask:
model = DataParallelModel.load(
UNet(downsample=5, out_channels=6).cuda(), model_file)
elif model_file is not None:
print('here')
#model = DataParallelModel.load(UNet(downsample=5).cuda(), model_file)
model = DataParallelModel.load(UNet(downsample=6).cuda(), model_file)
else:
model = Transfer(src_task=tasks.rgb, dest_task=dest_task).load_model()
model.compile(torch.optim.Adam, lr=3e-4, weight_decay=2e-6, amsgrad=True)
# DATA LOADING 1
results = model.predict(data)[:, -3:].clamp(min=0, max=1)
if results.shape[1] == 1:
results = torch.cat([results] * 3, dim=1)
if percep_mode:
percep_model = Transfer(src_task=dest_task,
dest_task=tasks.normal).load_model()
percep_model.eval()
eval_loader = torch.utils.data.DataLoader(
torch.utils.data.TensorDataset(results),
batch_size=16,
num_workers=16,
shuffle=False,
pin_memory=True)
final_preds = []
for preds, in eval_loader:
print('preds shape', preds.shape)
final_preds += [percep_model.forward(preds[:, -3:])]
results = torch.cat(final_preds, dim=0)
return results
def run_eval_suite(name,
dest_task=tasks.normal,
graph_file=None,
model_file=None,
logger=None,
sample=800,
show_images=False,
old=False):
if graph_file is not None:
graph = TaskGraph(tasks=[tasks.rgb, dest_task], pretrained=False)
graph.load_weights(graph_file)
model = graph.edge(tasks.rgb, dest_task).load_model()
elif old:
model = DataParallelModel.load(UNetOld().cuda(), model_file)
elif model_file is not None:
#model = DataParallelModel.load(UNet(downsample=5).cuda(), model_file)
model = DataParallelModel.load(UNet(downsample=6).cuda(), model_file)
else:
model = Transfer(src_task=tasks.normal,
dest_task=dest_task).load_model()
model.compile(torch.optim.Adam, lr=3e-4, weight_decay=2e-6, amsgrad=True)
dataset = ValidationMetrics("almena", dest_task=dest_task)
result = dataset.evaluate(model, sample=800)
logger.text(name + ": " + str(result))
def run_perceptual_eval_suite(name,
intermediate_task=tasks.normal,
dest_task=tasks.normal,
graph_file=None,
model_file=None,
logger=None,
sample=800,
show_images=False,
old=False,
perceptual_transfer=None,
multitask=False):
if perceptual_transfer is None:
percep_model = Transfer(src_task=intermediate_task,
dest_task=dest_task).load_model()
if graph_file is not None:
graph = TaskGraph(tasks=[tasks.rgb, intermediate_task],
pretrained=False)
graph.load_weights(graph_file)
model = graph.edge(tasks.rgb, intermediate_task).load_model()
elif old:
model = DataParallelModel.load(UNetOld().cuda(), model_file)
elif multitask:
print('running multitask')
model = DataParallelModel.load(
UNet(downsample=5, out_channels=6).cuda(), model_file)
elif model_file is not None:
#model = DataParallelModel.load(UNet(downsample=5).cuda(), model_file)
model = DataParallelModel.load(UNet(downsample=6).cuda(), model_file)
else:
model = Transfer(src_task=tasks.rgb,
dest_task=intermediate_task).load_model()
model.compile(torch.optim.Adam, lr=3e-4, weight_decay=2e-6, amsgrad=True)
dataset = ValidationMetrics("almena", dest_task=dest_task)
result = dataset.evaluate_with_percep(model,
sample=800,
percep_model=percep_model)
logger.text(name + ": " + str(result))
def main(loss_config="gt_mse", mode="standard", pretrained=False, batch_size=64, **kwargs):
# MODEL
# model = DataParallelModel.load(UNet().cuda(), "standardval_rgb2normal_baseline.pth")
model = functional_transfers.n.load_model() if pretrained else DataParallelModel(UNet())
model.compile(torch.optim.Adam, lr=(3e-5 if pretrained else 3e-4), weight_decay=2e-6, amsgrad=True)
scheduler = MultiStepLR(model.optimizer, milestones=[5*i + 1 for i in range(0, 80)], gamma=0.95)
# FUNCTIONAL LOSS
functional = get_functional_loss(config=loss_config, mode=mode, model=model, **kwargs)
print (functional)
# LOGGING
logger = VisdomLogger("train", env=JOB)
logger.add_hook(lambda logger, data: logger.step(), feature="loss", freq=20)
logger.add_hook(lambda logger, data: model.save(f"{RESULTS_DIR}/model.pth"), feature="loss", freq=400)
logger.add_hook(lambda logger, data: scheduler.step(), feature="epoch", freq=1)
functional.logger_hooks(loggers)
# DATA LOADING
ood_images = load_ood(ood_path=f'{BASE_DIR}/data/ood_images/')
train_loader, val_loader, train_step, val_step = load_train_val([tasks.rgb, tasks.normal], batch_size=batch_size)
# train_buildings=["almena"], val_buildings=["almena"])
test_set, test_images = load_test("rgb", "normal")
logger.images(test_images, "images", resize=128)
logger.images(torch.cat(ood_images, dim=0), "ood_images", resize=128)
# TRAINING
for epochs in range(0, 800):
preds_name = "start_preds" if epochs == 0 and pretrained else "preds"
ood_name = "start_ood" if epochs == 0 and pretrained else "ood"
plot_images(model, logger, test_set, dest_task="normal", ood_images=ood_images,
loss_models=functional.plot_losses, preds_name=preds_name, ood_name=ood_name
)
logger.update("epoch", epochs)
logger.step()
train_set = itertools.islice(train_loader, train_step)
val_set = itertools.islice(val_loader, val_step)
val_metrics = model.predict_with_metrics(val_set, loss_fn=functional, logger=logger)
train_metrics = model.fit_with_metrics(train_set, loss_fn=functional, logger=logger)
functional.logger_update(logger, train_metrics, val_metrics)
def __init__(self, cfg):
super(MONet, self).__init__()
# Configuration
self.K_steps = cfg.K_steps
self.prior_mode = cfg.prior_mode
self.mckl = cfg.montecarlo_kl
self.debug = cfg.debug
self.pixel_bound = cfg.pixel_bound
# Sub-Modules
# - Attention Network
if not hasattr(cfg, 'filter_start'):
cfg['filter_start'] = 32
core = UNet(int(np.log2(cfg.img_size) - 1), cfg.filter_start)
self.att_process = seq_att.SimpleSBP(core)
# - Component VAE
self.comp_vae = ComponentVAE(nout=4, cfg=cfg, act=nn.ReLU())
self.comp_vae.pixel_bound = False
# Initialise pixel output standard deviations
std = cfg.pixel_std2 * torch.ones(1, 1, 1, 1, self.K_steps)
std[0, 0, 0, 0, 0] = cfg.pixel_std1 # first step
self.register_buffer('std', std)
def get_model(src_task, dest_task):
if isinstance(src_task, str) and isinstance(dest_task, str):
src_task, dest_task = get_task(src_task), get_task(dest_task)
if (src_task.name, dest_task.name) in model_types:
return model_types[(src_task.name, dest_task.name)]()
elif isinstance(src_task, ImageTask) and isinstance(dest_task, ImageTask):
return UNet(downsample=5,
in_channels=src_task.shape[0],
out_channels=dest_task.shape[0])
elif isinstance(src_task, ImageTask) and isinstance(dest_task, ClassTask):
return ResNet(in_channels=src_task.shape[0],
out_channels=dest_task.classes)
elif isinstance(src_task, ImageTask) and isinstance(
dest_task, PointInfoTask):
return ResNet(out_channels=dest_task.out_channels)
return None
def __call__(self, im):
return im.filter(self.filter)
def __repr__(self):
return 'GaussianBulr Filter with Radius {:d}'.format(self.radius)
""" Model definitions for launching new transfer jobs between tasks. """
model_types = {
('normal', 'principal_curvature'):
lambda: Dense1by1Net(),
('normal', 'depth_zbuffer'):
lambda: UNetDepth(),
('normal', 'reshading'):
lambda: UNet(downsample=5),
('depth_zbuffer', 'normal'):
lambda: UNet(downsample=6, in_channels=1, out_channels=3),
('reshading', 'normal'):
lambda: UNet(downsample=4, in_channels=3, out_channels=3),
('sobel_edges', 'principal_curvature'):
lambda: UNet(downsample=5, in_channels=1, out_channels=3),
('depth_zbuffer', 'principal_curvature'):
lambda: UNet(downsample=4, in_channels=1, out_channels=3),
('principal_curvature', 'depth_zbuffer'):
lambda: UNet(downsample=6, in_channels=3, out_channels=1),
('rgb', 'normal'):
lambda: UNet(downsample=6),
('rgb', 'keypoints2d'):
lambda: UNet(downsample=3, out_channels=1),
}
from modules.percep_nets import DenseNet, Dense1by1Net, DenseKernelsNet, DeepNet, BaseNet, WideNet, PyramidNet
from modules.depth_nets import UNetDepth
from modules.unet import UNet, UNetOld, UNetOld2, UNetReshade
from modules.resnet import ResNetClass
from fire import Fire
import IPython
pretrained_transfers = {
('normal', 'principal_curvature'): (lambda: Dense1by1Net(),
f"{MODELS_DIR}/normal2curvature.pth"),
('normal', 'depth_zbuffer'): (lambda: UNetDepth(),
f"{MODELS_DIR}/normal2zdepth_zbuffer.pth"),
('normal', 'sobel_edges'):
(lambda: UNet(out_channels=1, downsample=4).cuda(),
f"{MODELS_DIR}/normal2edges2d.pth"),
('normal', 'reshading'): (lambda: UNetReshade(downsample=5),
f"{MODELS_DIR}/normal2reshade.pth"),
('normal', 'keypoints3d'): (lambda: UNet(downsample=5, out_channels=1),
f"{MODELS_DIR}/normal2keypoints3d.pth"),
('normal', 'keypoints2d'): (lambda: UNet(downsample=5, out_channels=1),
f"{MODELS_DIR}/normal2keypoints2d_new.pth"),
('normal', 'edge_occlusion'): (lambda: UNet(downsample=5, out_channels=1),
f"{MODELS_DIR}/normal2edge_occlusion.pth"),
('depth_zbuffer', 'normal'): (lambda: UNet(in_channels=1, downsample=6),
f"{MODELS_DIR}/depth2normal.pth"),
('depth_zbuffer', 'sobel_edges'):
(lambda: UNet(downsample=4, in_channels=1, out_channels=1).cuda(),
f"{MODELS_DIR}/depth_zbuffer2sobel_edges.pth"),
('depth_zbuffer', 'principal_curvature'):
def run_viz_suite(name,
data_loader,
dest_task=tasks.depth_zbuffer,
graph_file=None,
model_file=None,
old=False,
multitask=False,
percep_mode=None,
downsample=6,
out_channels=3,
final_task=tasks.normal,
oldpercep=False):
extra_task = [final_task] if percep_mode else []
if graph_file is not None:
graph = TaskGraph(tasks=[tasks.rgb, dest_task] + extra_task,
pretrained=False)
graph.load_weights(graph_file)
model = graph.edge(tasks.rgb, dest_task).load_model()
elif old:
model = DataParallelModel.load(UNetOld().cuda(), model_file)
elif multitask:
model = DataParallelModel.load(
UNet(downsample=5, out_channels=6).cuda(), model_file)
elif model_file is not None:
# downsample = 5 or 6
print('loading main model')
#model = DataParallelModel.load(UNetReshade(downsample=downsample, out_channels=out_channels).cuda(), model_file)
model = DataParallelModel.load(
UNet(downsample=downsample, out_channels=out_channels).cuda(),
model_file)
#model = DataParallelModel.load(UNet(downsample=6).cuda(), model_file)
else:
model = DummyModel(
Transfer(src_task=tasks.rgb, dest_task=dest_task).load_model())
model.compile(torch.optim.Adam, lr=3e-4, weight_decay=2e-6, amsgrad=True)
# DATA LOADING 1
results = []
final_preds = []
if percep_mode:
print('Loading percep model...')
if graph_file is not None and not oldpercep:
percep_model = graph.edge(dest_task, final_task).load_model()
percep_model.compile(torch.optim.Adam,
lr=3e-4,
weight_decay=2e-6,
amsgrad=True)
else:
percep_model = Transfer(src_task=dest_task,
dest_task=final_task).load_model()
percep_model.eval()
print("Converting...")
for data, in data_loader:
preds = model.predict_on_batch(data)[:, -3:].clamp(min=0, max=1)
results.append(preds.detach().cpu())
if percep_mode:
try:
final_preds += [
percep_model.forward(preds[:, -3:]).detach().cpu()
]
except RuntimeError:
preds = torch.cat([preds] * 3, dim=1)
final_preds += [
percep_model.forward(preds[:, -3:]).detach().cpu()
]
#break
if percep_mode:
results = torch.cat(final_preds, dim=0)
else:
results = torch.cat(results, dim=0)
return results
# get target task and model
target_tasks = ['normal', 'depth_zbuffer', 'reshading']
try:
task_index = target_tasks.index(args.task)
target_task = target_tasks[task_index]
except:
print(
"task should be one of the following: normal, depth_zbuffer, reshading"
)
sys.exit()
#models = [UNet(), UNet(downsample=6, out_channels=1), UNetReshade(downsample=5)]
#direct, emboss4d, grey, laplace, gauss, sobel, wav, sharp
models_normal = [
UNet(out_channels=6),
UNet(in_channels=4, downsample=6, out_channels=6),
UNet(downsample=6, in_channels=1, out_channels=6),
UNet(in_channels=1, downsample=6, out_channels=6),
UNet(in_channels=3, downsample=6, out_channels=6),
UNet(downsample=5, in_channels=1, out_channels=6),
UNet(downsample=6, in_channels=12, out_channels=6),
UNet(in_channels=3, downsample=6, out_channels=6)
]
models_reshading = [
UNet(downsample=5, out_channels=2),
UNet(in_channels=4, downsample=5, out_channels=2),
UNet(downsample=5, in_channels=1, out_channels=2),
UNet(in_channels=1, downsample=5, out_channels=2),
UNet(in_channels=3, downsample=5, out_channels=2),
UNet(downsample=5, in_channels=1, out_channels=2),
from modules.resnet import ResNetClass
from fire import Fire
import IPython
###
pretrained_transfers = {
# percep models used in consistency training
('normal', 'principal_curvature'):
(lambda: Dense1by1Net(), f"{MODELS_DIR}/perceps/normal2curvature.pth"),
('normal', 'depth_zbuffer'):
(lambda: UNetDepth(), f"{MODELS_DIR}/perceps/normal2zdepth_zbuffer.pth"),
('normal', 'sobel_edges'):
(lambda: UNet(out_channels=1, downsample=4).cuda(),
f"{MODELS_DIR}/perceps/normal2edges2d.pth"),
('normal', 'reshading'): (lambda: UNetReshade(downsample=5),
f"{MODELS_DIR}/perceps/normal2reshade.pth"),
('normal', 'keypoints3d'):
(lambda: UNet(downsample=5, out_channels=1),
f"{MODELS_DIR}/perceps/normal2keypoints3d.pth"),
('normal', 'keypoints2d'):
(lambda: UNet(downsample=5, out_channels=1),
f"{MODELS_DIR}/perceps/normal2keypoints2d.pth"),
('normal', 'edge_occlusion'):
(lambda: UNet(downsample=5, out_channels=1),
f"{MODELS_DIR}/perceps/normal2edge_occlusion.pth"),
('depth_zbuffer', 'sobel_edges'):
(lambda: UNet(downsample=4, in_channels=1, out_channels=1).cuda(),
f"{MODELS_DIR}/perceps/depth_zbuffer2sobel_edges.pth"),
from fire import Fire
import IPython
pretrained_transfers = {
(('depth_zbuffer', 'FoV', 'normal'), 'normal'):
(lambda: GeoNetDepthToNormal(),
f"{MODELS_DIR}/geonet_depth_to_normal2.pth"),
(('depth_zbuffer', 'FoV', 'normal'), 'depth_zbuffer'):
(lambda: GeoNetNormalToDepth(), None),
('normal', 'principal_curvature'):
(lambda: Dense1by1Net(), f"{MODELS_DIR}/normal2curvature_dense_1x1.pth"),
('normal', 'depth_zbuffer'): (lambda: UNetDepth(),
f"{MODELS_DIR}/normal2zdepth_unet_v4.pth"),
('normal', 'sobel_edges'):
(lambda: UNet(out_channels=1, downsample=4).cuda(),
f"{MODELS_DIR}/normal2edges2d_sobel_unet4.pth"),
('sobel_edges', 'normal'):
(lambda: UNet(in_channels=1, downsample=5).cuda(),
f"{MODELS_DIR}/sobel_edges2normal.pth"),
('normal', 'grayscale'):
(lambda: UNet(out_channels=1, downsample=6).cuda(),
f"{MODELS_DIR}/normals2gray_unet.pth"),
('principal_curvature', 'normal'):
(lambda: UNetOld2(),
f"{MODELS_DIR}/results_inverse_cycle_unet1x1model.pth"),
('principal_curvature', 'sobel_edges'):
(lambda: UNet(downsample=4, out_channels=1),
f"{MODELS_DIR}/principal_curvature2sobel_edges.pth"),
# ('depth_zbuffer', 'normal'):
# (lambda: UNet(in_channels=1, downsample=4), f"{MODELS_DIR}/depth2normal_unet4.pth"),
from modules.unet import UNet
from modules.percep_nets import ResidualsNetDown, ResidualsNetUp
from task_configs import get_task
from utils import *
model_types = {
'unet_based': {
'normal': {
'down':
(lambda: UNet(in_channels=3, downsample=3), f"normal2LS.pth"),
'up':
(lambda: UNet(out_channels=3, downsample=3), f"LS2normal.pth"),
},
'sobel_edges': {
'down':
(lambda: UNet(downsample=3, in_channels=1), f"sobel_edges2LS.pth"),
'up': (lambda: UNet(out_channels=1, downsample=3),
f"LS2sobel_edges.pth"),
},
'reshading': {
'down':
(lambda: UNet(downsample=3, in_channels=3), f"reshading2LS.pth"),
'up':
(lambda: UNet(downsample=3, out_channels=3), f"LS2reshading.pth"),
},
'keypoints2d': {
'down':
(lambda: UNet(downsample=3, in_channels=1), f"keypoints2d2LS.pth"),
'up': (lambda: UNet(downsample=3, out_channels=1),
f"LS2keypoints2d.pth"),
},
transforms.CenterCrop(256),
transforms.ToTensor()
])
trans_topil = transforms.ToPILImage()
os.system(f"mkdir -p {args.output_path}")
# get target task and model
target_tasks = ['normal', 'depth', 'reshading']
try:
task_index = target_tasks.index(args.task)
except:
print("task should be one of the following: normal, depth, reshading")
sys.exit()
models = [
UNet(),
UNet(downsample=6, out_channels=1),
UNetReshade(downsample=5)
]
model = models[task_index]
map_location = (lambda storage, loc: storage.cuda()
) if torch.cuda.is_available() else torch.device('cpu')
def save_outputs(img_path, output_file_name):
img = Image.open(img_path)
img_tensor = trans_totensor(img)[:3].unsqueeze(0)
# compute baseline and consistency output