1 files changed, 36 insertions, 0 deletions
diff --git a/precompute_feature_maps.py b/precompute_feature_maps.py
new file mode 100755
index 0000000..a631b9c
--- /dev/null
+++ b/precompute_feature_maps.py
@@ -0,0 +1,36 @@
+### Copyright (C) 2017 NVIDIA Corporation. All rights reserved. 
+### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+from options.train_options import TrainOptions
+from data.data_loader import CreateDataLoader
+from models.models import create_model
+import numpy as np
+import os, time
+import util.util as util
+from torch.autograd import Variable
+import torch.nn as nn
+
+opt = TrainOptions().parse()
+opt.nThreads = 1
+opt.batchSize = 1 
+opt.serial_batches = True 
+opt.no_flip = True
+opt.instance_feat = True
+
+name = 'features'
+save_path = os.path.join(opt.checkpoints_dir, opt.name)
+
+############ Initialize #########
+data_loader = CreateDataLoader(opt)
+dataset = data_loader.load_data()
+dataset_size = len(data_loader)
+model = create_model(opt)
+util.mkdirs(os.path.join(opt.dataroot, opt.phase + '_feat'))
+
+######## Save precomputed feature maps for 1024p training #######
+for i, data in enumerate(dataset):
+	print('%d / %d images' % (i+1, dataset_size)) 
+	feat_map = model.module.netE.forward(Variable(data['image'].cuda(), volatile=True), data['inst'].cuda())
+	feat_map = nn.Upsample(scale_factor=2, mode='nearest')(feat_map)
+	image_numpy = util.tensor2im(feat_map.data[0])
+	save_path = data['path'][0].replace('/train_label/', '/train_feat/')
+	util.save_image(image_numpy, save_path)
+\ No newline at end of file