diff options
| -rw-r--r-- | data/recursive_dataset.py | 2 | ||||
| -rw-r--r-- | options/test_options.py | 4 | ||||
| -rw-r--r-- | test.py | 24 |
3 files changed, 20 insertions, 10 deletions
diff --git a/data/recursive_dataset.py b/data/recursive_dataset.py index b51777f..4bd0798 100644 --- a/data/recursive_dataset.py +++ b/data/recursive_dataset.py @@ -35,7 +35,7 @@ class RecursiveDataset(BaseDataset): self.transform = get_transform(opt) def __getitem__(self, index): - A_path = "/home/lens/code/pytorch-CycleGAN-and-pix2pix/recursive/frame_{:04d}.png".format(index) + A_path = self.opt.results_dir + "/" + self.opt.name + "/exp:" + self.opt.experiment + "/frame_{:05d}.png".format(index) print('next image: ' + A_path) A_img = Image.open(A_path).convert('RGB') A = self.transform(A_img) diff --git a/options/test_options.py b/options/test_options.py index 6b79860..f8c3a31 100644 --- a/options/test_options.py +++ b/options/test_options.py @@ -1,9 +1,11 @@ from .base_options import BaseOptions - +import random class TestOptions(BaseOptions): def initialize(self): BaseOptions.initialize(self) + self.parser.add_argument('--start_img', type=str, default="test.jpg", help='starting image for recursive generation') + self.parser.add_argument('--experiment', type=str, default="exp_{:06d}".format(random.randint(1,1000001)), help='experiment name (name of directory)') self.parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.') self.parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.') self.parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images') @@ -3,17 +3,24 @@ from options.test_options import TestOptions from data import CreateRecursiveDataLoader from models import create_model from util.visualizer import Visualizer -from util import html - +from util import html, mkdir +from shutil import movefile import subprocess + if __name__ == '__main__': opt = TestOptions().parse() 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 + opt.experiment = opt.start_img.split(".")[0] + + render_dir = self.opt.results_dir + "/" + self.opt.name + "/exp:" + opt.experiment + + mkdir(render_dir) + movefile(opt.start_img, render_dir) data_loader = CreateRecursiveDataLoader(opt) dataset = data_loader.load_data() @@ -35,11 +42,12 @@ if __name__ == '__main__': print('%04d: process image... %s' % (i, img_path)) ims = visualizer.save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio) if dataset.name() == 'RecursiveDatasetDataLoader': - last_image = "results/woodscaled_4_pix2pix/test_latest/images/" + ims[1] - next_image = "recursive/frame_{:04d}.png".format(i+1) - cmd = ("/bin/cp", last_image, next_image) - process = subprocess.Popen(cmd, stdout=subprocess.PIPE) - output, error = process.communicate() - dataset.append(next_image) + last_image = opt.results_dir + "/" + opt.name + "/test_latest/images/" + ims[1] + next_image = opt.results_dir + "/" + opt.name + "/exp:" + opt.experiment + "/frame_{:05d}.png".format(i+1) + movefile(last_image, next_image) webpage.save() + + cmd = ("/usr/bin/ffmpeg", "-i", "frame_%04d.png", "-c:v", "libx264", "-vf", "fps=30", "-pix_fmt", "yuv420p", opt.name + "_" + opt.experiment + ".mp4") + process = subprocess.Popen(cmd, stdout=subprocess.PIPE) + output, error = process.communicate() |