diff options
| -rwxr-xr-x | .gitignore | 4 | ||||
| -rw-r--r-- | options/dataset_options.py | 374 | ||||
| -rw-r--r-- | recursive.py | 37 | ||||
| -rw-r--r-- | recursive.sh | 6 |
4 files changed, 409 insertions, 12 deletions
@@ -1,4 +1,7 @@ debug* +datasets/ +sequences/ +a_b/ checkpoints/ results/ build/ @@ -38,3 +41,4 @@ test/.coverage test/data/legacy_serialized.pt *.DS_Store *~ + diff --git a/options/dataset_options.py b/options/dataset_options.py new file mode 100644 index 0000000..fdea637 --- /dev/null +++ b/options/dataset_options.py @@ -0,0 +1,374 @@ +from .base_options import BaseOptions + +class DatasetOptions(BaseOptions): + def initialize(self, args=None): + # BaseOptions.initialize(self) + # type = int, float, str OR action='store_true' + + # self.parser.add_argument( + # '--', + # type=int, + # default=0, + # help='' + # ) + + required = args is None + + self.parser.add_argument( + '--in_dir', + type=str, + required=required, + help='input directory' + ) + + self.parser.add_argument( + '--out_dir', + type=str, + required=required, + help='output directory' + ) + + self.parser.add_argument( + '--split', + action='store_true', + help='construct train/test/split for this output, as A' + ) + + self.parser.add_argument( + '--ab', + action='store_true', + help='construct test split into directory A and B, where B = in_dir' + ) + + self.parser.add_argument( + '--mov', + action='store_true', + help='generate video from output directory' + ) + + self.parser.add_argument( + '--render-frames', + action='store_true', + help='render the source frames only' + ) + + self.parser.add_argument( + '--tag', + type=str, + default="", + help='another way to tag this whole thing' + ) + + self.parser.add_argument( + '--scp', + action='store_true', + help='scp this file somewhere' + ) + + self.parser.add_argument( + '--scp-to', + type=str, + default="jules@asdf.us:asdf/neural/", + help='scp destination' + ) + + ## LIVE IMAGE PROCESSING + + self.parser.add_argument( + '--send-image', + type=str, + default='b', + help='which image to send... a, b, recursive, sequence' + ) + + self.parser.add_argument( + '--store-a', + action='store_true', + help='dont remove the generated A image after processing it' + ) + + self.parser.add_argument( + '--store-b', + action='store_true', + help='after generating a B image, save it to disk' + ) + + self.parser.add_argument( + '--exit', + action='store_true', + help='exit immediately if set to true (used to interrupt)' + ) + + self.parser.add_argument( + '--output-format', + type=str, + default="JPEG", + help='what gets sent over the wire - JPEG or PNG' + ) + + self.parser.add_argument( + '--frame-delay', + type=float, + default=1.0, + help='delay per render, in seconds. inverse of frame rate' + ) + + ## LOAD A NEW SEQUENCE + + self.parser.add_argument( + '--load-sequence', + action='store_true', + help='instruction to reload the sequence after training the next image' + ) + + self.parser.add_argument( + '--sequence-name', + type=str, + default='woodscaled_4', + help='which sequence to load' + ) + + self.parser.add_argument( + '--seek_to', + type=int, + default=1, + help='initial offset into sequence' + ) + + ## LOAD A NEW CHECKPOINT + + self.parser.add_argument( + '--load-checkpoint', + action='store_true', + help='instruction to reload the checkpoint after training the next image' + ) + + self.parser.add_argument( + '--checkpoint-name', + type=str, + default='messi', + help='which checkpoint to load' + ) + + self.parser.add_argument( + '--epoch', + type=str, + default='latest', + help='which epoch to load' + ) + + ## IMAGE FILTERS + + ### RECURSION + + self.parser.add_argument( + '--processing', + action='store_true', + help='internally stores whether or not we should be processing..' + ) + + self.parser.add_argument( + '--pause', + action='store_true', + help='internally tells the loop to pause..' + ) + + self.parser.add_argument( + '--transition', + action='store_true', + help=' sine wave transition' + ) + + self.parser.add_argument( + '--transition-period', + default=3000, + type=int, + help='period of sine wave transition' + ) + + self.parser.add_argument( + '--transition-min', + default=1e-3, + type=float, + help='minimum amount of stabilization to apply' + ) + + self.parser.add_argument( + '--transition-max', + default=1.0, + type=float, + help='maximum amount of stabilization to apply' + ) + + self.parser.add_argument( + '--recursive', + action='store_true', + help='recurse on previous output' + ) + + self.parser.add_argument( + '--recursive-frac', + default=0.3, + type=float, + help='amount of previous step to use in recursion' + ) + + self.parser.add_argument( + '--just-copy', + action='store_true', + help='dont preprocess first frame', + ) + + self.parser.add_argument( + '--sequence', + action='store_true', + help='recurse guided by image sequence' + ) + + self.parser.add_argument( + '--sequence-frac', + default=0.2, + type=float, + help='amount of sequence image to use in recursion' + ) + + self.parser.add_argument( + '--recurse-roll', + default=0, + type=int, + help='in px, roll mixed recursed image 1px to reduce vertical feedback' + ) + + self.parser.add_argument( + '--recurse-roll-axis', + default=1, + type=int, + help='axis of roll' + ) + + self.parser.add_argument( + '--process-frac', + default=0.5, + type=float, + help='amount of processed image (clahe, poster, etc) to use in feeder step' + ) + + ### GRAYSCALE + + self.parser.add_argument( + '--grayscale', + action='store_true', + help='convert image to grayscale first' + ) + + ### CLAHE + + self.parser.add_argument( + '--clahe', + action='store_true', + help='apply clahe contrast correction' + ) + + self.parser.add_argument( + '--clip-limit', + default=2.0, + type=float, + help='clip limit for clahe algorithm (1.0 is subtle, 4.0 is aggressive)' + ) + + ### POSTERIZE + + self.parser.add_argument( + '--posterize', + action='store_true', + help='posterize image' + ) + + self.parser.add_argument( + '--spatial-window', + default=16, + type=int, + help='spatial window for quantize' + ) + + self.parser.add_argument( + '--color-window', + default=64, + type=int, + help='color window for quantize' + ) + + ### BRIGHTNESS GRADIENT + + self.parser.add_argument( + '--brightness-gradient', + action='store_true', + help='gradiate from first frame to last (done in Lab colorspace)' + ) + + self.parser.add_argument( + '--brightness-sigma', + default=64, + type=int, + help='width of brightness gradient along L axis' + ) + + ### BLUR + + self.parser.add_argument( + '--blur', + action='store_true', + help='blur image by N' + ) + + self.parser.add_argument( + '--blur-radius', + default=3, + type=int, + help='blur sigma' + ) + + self.parser.add_argument( + '--blur-sigma', + default=0.0, + type=float, + help='blur sigma' + ) + + ### CANNY EDGE DETECTION + + self.parser.add_argument( + '--canny', + action='store_true', + help='do canny edge detection on image' + ) + + self.parser.add_argument( + '--canny-lo', + default=100, + type=int, + help='canny low threshold' + ) + + self.parser.add_argument( + '--canny-hi', + default=200, + type=int, + help='canny high threshold' + ) + + def parse(self, args=None): + if not self.initialized: + self.initialize(args) + if args is not None: + self.opt, unknown = self.parser.parse_known_args(args) + else: + self.opt, unknown = self.parser.parse_known_args() + + argz = vars(self.opt) + print('------------ Options -------------') + for k, v in sorted(argz.items()): + print('%s: %s' % (str(k), str(v))) + print('-------------- End ----------------') + + return self.opt
\ No newline at end of file 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() diff --git a/recursive.sh b/recursive.sh index 667fa8e..df1bf93 100644 --- a/recursive.sh +++ b/recursive.sh @@ -1,7 +1,9 @@ #!/bin/bash if [ "$1" == "" ]; then - echo "Usage: $0 [dataset]" + echo "Usage: $0 [dataset] [start_file]" exit 1 fi -python recursive.py --dataroot "./datasets/${1}/" --name "$1" --label_nc 0 --no_instance --which_epoch latest +d=`date +%s` +mkdir -p "./results/${1}_${d}" +python recursive.py --dataroot "./results/${1}/" --start_file "${2}" --name "$1" --label_nc 0 --no_instance --which_epoch latest --how_many 2000 |