DATASET OPTIONZ"

1 files changed, 27 insertions, 10 deletions

diff --git a/recursive.py b/recursive.py
index dc08b28..da32a8a 100644
--- a/recursive.py
+++ b/recursive.py
@@ -3,6 +3,7 @@
 import os
 from collections import OrderedDict
 from options.test_options import TestOptions
+from options.dataset_options import DatasetOptions
 from data.data_loader import CreateDataLoader
 from models.models import create_model
 import util.util as util
@@ -12,17 +13,28 @@ import torch
 from run_engine import run_trt_engine, run_onnx
 
 opt = TestOptions().parse(save=False)
+data_opt = DatasetOptions().parse(opt.unknown)
 opt.nThreads = 1   # test code only supports nThreads = 1
 opt.batchSize = 1  # test code only supports batchSize = 1
 opt.serial_batches = True  # no shuffle
 opt.no_flip = True  # no flip
+if data_opt.tag == '':
+    d = datetime.now()
+    tag = data_opt.tag = "{}_{}_{}".format(
+        opt.name,
+        opt.experiment,
+        d.strftime('%Y%m%d%H%M')
+    )
+else:
+    tag = data_opt.tag
+
+opt.render_dir = render_dir = opt.results_dir + opt.name + "/" + tag + "/"
 
 data_loader = CreateDataLoader(opt)
 dataset = data_loader.load_data()
-visualizer = Visualizer(opt)
-# create website
-web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
-webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
+
+start_img_path = os.path.join(render_dir, "frame_00000.png")
+copyfile(opt.start_img, start_img_path)
 
 for i, data in enumerate(dataset):
     if i >= opt.how_many:
@@ -36,10 +48,15 @@ for i, data in enumerate(dataset):
     minibatch = 1 
     generated = model.inference(data['label'], data['inst'])
         
-    visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)),
-                           ('synthesized_image', util.tensor2im(generated.data[0]))])
-    img_path = data['path']
-    print('process image... %s' % img_path)
-    visualizer.save_images(webpage, visuals, img_path)
+    last_path = opt.render_dir + "frame_{:05d}.png".format(i)
+    tmp_path = opt.render_dir + "frame_{:05d}_tmp.png".format(i+1)
+    next_path = opt.render_dir + "frame_{:05d}.png".format(i+1)
+    current_path = opt.render_dir + "ren_{:05d}.png".format(i+1)
+    print('process image... %s' % last_path)
+
+    im = util.tensor2im(generated.data[0])
+    image_pil = Image.fromarray(im, mode='RGB')
+    image_pil.save(tmp_path)
+    os.rename(tmp_path, next_path)
+
 
-webpage.save()