示例#1
0
def test(contentpath, stylepath, multi=False):
    '''一次前传得到风格化图像'''
    if multi == False:
        content_name = (contentpath.split('/')[-1]).split('.')[0]
        style_name = (stylepath.split('/')[-1]).split('.')[0]
    else:
        content_name = (contentpath.split('\\')[-1]).split('.')[0]
        style_name = (stylepath.split('\\')[-1]).split('.')[0]

    transfer = test_transform(512)

    contentimg = Image.open(str(contentpath)).convert('RGB')
    styleimg = Image.open(str(stylepath)).convert('RGB')
    if args.preserve_color: styleimg = change_color(styleimg, contentimg)
    contentimg = transfer(contentimg).unsqueeze(0)
    styleimg = transfer(styleimg).unsqueeze(0)

    #if args.preserve_color: styleimg = coral(styleimg, contentimg)

    decoder = Decoder().to(device).eval()
    decoder.load_state_dict(torch.load(args.model_path))

    fbnet = FPnet(decoder, True).to(device).eval()
    output = fbnet(contentimg,
                   styleimg,
                   alpha=args.alpha,
                   lamda=args.lamda,
                   require_loss=False)

    image_name = args.save_dir + '/' + content_name + '+' + style_name + '.jpg'
    save_image(output.cpu(), image_name)
    print('image saved  as:  ' + image_name)
    contentimg.detach()
    styleimg.detach()
示例#2
0
def main():
    parser = argparse.ArgumentParser(description='chainer implementation of pix2pix')
    parser.add_argument('--gpu', '-g', type=int, default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--img', '-i', help='Input image')
    parser.add_argument('--out', '-o', default='result_dehighlight',
                        help='Directory to output the result')
    args = parser.parse_args()


    # Set up a neural network to train
    enc = Encoder(in_ch=3)
    dec = Decoder(out_ch=1)

    chainer.serializers.load_npz(ENC_W, enc)
    chainer.serializers.load_npz(DEC_W, dec)

    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()

    inimg = loadimg(args.img)
    ch, h, w = inimg.shape
    # add paddings so that input array has the size of mutiples of 256.
    in_ary = np.zeros((ch,math.ceil(h/256)*256, math.ceil(w/256)*256), dtype="f")
    in_ary[:,0:h,0:w] = inimg
    x_in = in_ary[np.newaxis,:] # to fit into the minibatch shape
    print(x_in.shape)
    # x_in as an input image
    x_in = chainer.Variable(x_in)
    if args.gpu >= 0:
        x_in.to_gpu()

    st = time.time()

    for i in range(10):

        z = enc(x_in)
        x_out = dec(z)

    ts = time.time() - st
    print("time:{:.2f}".format(ts/10.0))


    if args.gpu >= 0:
        out_ary = x_out.data.get()[0]
    else:
        out_ary = x_out.data[0]
    #img_show = np.zeros((inimg.shape[0], inimg.shape[1], inimg.shape[2]*2))
    #img_show[:,:,:inimg.shape[2]] = inimg
    #img_show[:,:outimg.shape[1],inimg.shape[2]:inimg.shape[2]+outimg.shape[2]] = outimg
    outimg = out_ary[:,0:h,0:w] # trim paddings
    img_show = np.concatenate((inimg, outimg), axis=2)
    bgrpic = to_bgr(img_show).copy()
    cv2.putText(bgrpic,"input",(3,15),cv2.FONT_HERSHEY_DUPLEX, 0.5,(255,0,0))
    cv2.putText(bgrpic,"output",(w+3,15),cv2.FONT_HERSHEY_DUPLEX, 0.5,(255,0,0))
    cv2.imshow("result", bgrpic)
    cv2.waitKey()
示例#3
0
def main():
    parser = argparse.ArgumentParser(description='chainer implementation of pix2pix')
    parser.add_argument('--gpu', '-g', type=int, default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed')
    parser.add_argument('--model', '-m', default='',
                        help='model snapshot')
    parser.add_argument('--enc', '-e', type=str, default='enc_iter_60000.npz', help='encoder snapshot')
    parser.add_argument('--dec', '-d', type=str, default='dec_iter_60000.npz', help='decoder snapshot')
    parser.add_argument('--out', '-o', type=str, default='out', help='output dir')
    parser.add_argument('--input', '-i', default='sample.jpg', help='input jpg', required=True)
    parser.add_argument('--contour', '-c', action='store_true', help='from contour image or not')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))

    enc = Encoder(in_ch=3)
    dec = Decoder(out_ch=3)
    dis = Discriminator(in_ch=3, out_ch=3)

    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
        dis.to_gpu()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer

    if args.model:
        opt_enc = make_optimizer(enc)
        opt_dec = make_optimizer(dec)
        opt_dis = make_optimizer(dis)

        # Set up a trainer
        updater = FacadeUpdater(
            models=(enc, dec, dis),
            iterator={},
            optimizer={
                'enc': opt_enc, 'dec': opt_dec,
                'dis': opt_dis},
            device=args.gpu)
        trainer = training.Trainer(updater, (200, 'epoch'), out='generate/')
        chainer.serializers.load_npz(args.model, trainer)
    elif args.enc and args.dec:
        chainer.serializers.load_npz(args.enc, enc)
        chainer.serializers.load_npz(args.dec, dec)

    if not args.contour:
        from make_contour import get_contour_image
        get_contour_image(args.input)

    generate_image_from_contour(args.input, enc, dec, args.out)
示例#4
0
def main():
    parser = argparse.ArgumentParser(
        description='chainer implementation of pix2pix')
    parser.add_argument('--gpu', '-g', type=int, default=-1,
        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--dec', '-d', type=str, default='dec_model.npz',
        help='decoder model')
    parser.add_argument('--enc', '-e', type=str, default='enc_model.npz',
        help='encoder model')
    args = parser.parse_args()

    enc = Encoder(in_ch=3)
    dec = Decoder(out_ch=1)

    chainer.serializers.load_npz(str(MODEL_PATH.joinpath(args.dec)), dec)
    chainer.serializers.load_npz(str(MODEL_PATH.joinpath(args.enc)), enc)

    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
    xp = enc.xp

    branch_images = np.zeros((12, 256, 256), dtype=np.uint8)
    plant_images = np.zeros((12, 256, 256, 3), dtype=np.uint8)
    p2b_images = np.zeros((12, 256, 256), dtype=np.uint8)

    for i in tqdm(range(1,12)):
        branch_path = DATASET_PATH.joinpath('branch', str(i))
        plant_path = DATASET_PATH.joinpath('plant', str(i))
        name = random.choice([_ for _ in branch_path.glob('*.png')]).name
        branch_image_path = branch_path.joinpath(name)
        plant_image_path = plant_path.joinpath(name)

        # open image
        branch_image = np.asarray(Image.open(branch_image_path).convert('L'))
        branch_images[i-1,:] = branch_image
        plant_image = np.asarray(Image.open(plant_image_path).convert('RGB'))
        plant_images[i-1,:] = plant_image

        plant_image = xp.array(plant_image).astype("f").transpose(2, 0, 1) / 128.0-1.0
        plant_image = plant_image.reshape(1, *plant_image.shape)

        with chainer.no_backprop_mode(), chainer.using_config('train', False):
            p2b_image = np.asarray(dec(enc(plant_image)).data.get())
            p2b_image = np.asarray(np.clip(p2b_image * 128 + 128, 0.0, 255.0), dtype=np.uint8).reshape(256, 256)
        p2b_images[i-1, :] = p2b_image

    Image.fromarray(branch_images.reshape(3, 4, 256, 256).transpose(0, 2, 1, 3).reshape(3*256, 4*256))\
        .save(str(RESULT_PATH.joinpath('branch_image.png')))
    Image.fromarray(plant_images.reshape(3, 4, 256, 256, 3).transpose(0, 2, 1, 3, 4).reshape(3*256, 4*256, 3))\
        .save(str(RESULT_PATH.joinpath('plant_image.png')))
    Image.fromarray(p2b_images.reshape(3, 4, 256, 256).transpose(0, 2, 1, 3).reshape(3*256, 4*256))\
        .save(str(RESULT_PATH.joinpath('p2b_image.png')))
示例#5
0
def test(contentpath,stylepath,pixel,multi=False):
    '''一次前传得到风格化图像'''
    if multi==False:
        content_name=(contentpath.split('/')[-1]).split('.')[0]
        style_name=(stylepath.split('/')[-1]).split('.')[0]
    else:
        content_name=(contentpath.split('\\')[-1]).split('.')[0]
        style_name=(stylepath.split('\\')[-1]).split('.')[0]

    mytransfer=test_transform(pixel)

    contentfile = open(str(contentpath),'rb')
    stylefile = open(str(stylepath),'rb')

    contentimg = Image.open(contentfile).convert('RGB')
    styleimg = Image.open(stylefile).convert('RGB')

    contentfile.close()
    stylefile.close()
    
    if args.preserve_color: styleimg = change_color(styleimg, contentimg)
    #if args.preserve_color: contentimg,styleimg,contentH,contentS = lumi_only(contentimg,styleimg,mytransfer)

    contentimg=mytransfer(contentimg).unsqueeze(0)
    styleimg=mytransfer(styleimg).unsqueeze(0)


    decoder=Decoder().to(device).eval()
    decoder.load_state_dict(torch.load(args.model_path))

    fbnet=FPnet(decoder,True).to(device).eval()
    output=fbnet(contentimg,styleimg,alpha=args.alpha,lamda=args.lamda,require_loss=False)

    #if args.preserve_color:output=recover_color(output,contentH,contentS)

    image_name=args.save_dir+'/'+content_name+'+'+style_name+'-'+str(pixel)+'-'+str(args.alpha)+'.jpg'

    save_image(output.cpu(),image_name)
    print('image saved  as:  '+image_name)

    contentimg.detach()
    styleimg.detach()
    output.detach()
示例#6
0
def style_transfer():
    enc = VGG19(input_shape)
    decs = [Decoder(enc, i, name='decoder %d' % i) for i in indices]
    sty_net = StyleTransfer(enc, decs)

    X = T.tensor4('input')
    Y = T.tensor4('style')

    mon = nn.Monitor(current_folder=args.path_to_weight_files)
    for idx, layer in enumerate(indices):
        weights = mon.load('decoder-%d-final.npz' % layer)
        nn.utils.numpy2shared(weights, decs[idx].params)

    nn.set_training_off()
    X_styled = sty_net(X, Y)
    transfer = nn.function([X, Y], X_styled, name='transfer style')

    if os.path.isfile(input_path_test) and os.path.isfile(style_path_test):
        input = prep_image_test(misc.imread(input_path_test))
        style = prep_image_test(misc.imread(style_path_test))
        output = transfer(input, style)
        mon.imwrite('test %s' % input_path_test[:-4], input)
        mon.imwrite('test %s' % style_path_test[:-4], style)
        mon.imwrite('test %s %s' % (input_path_test[:-4], style_path_test[:-4]), output)
    elif os.path.isfile(input_path_test) and os.path.isdir(style_path_test):
        input = prep_image_test(misc.imread(input_path_test))
        style_files = os.listdir(style_path_test)
        for style_file in style_files:
            style = prep_image_test(misc.imread(os.path.join(style_path_test, style_file)))
            output = transfer(input, style)
            mon.imwrite('test %s' % style_file[:-4], style)
            mon.imwrite('test %s %s' % (input_path_test[:-4], style_file[:-4]), output)
        mon.imwrite('test %s' % input_path_test[:-4], input)
    elif os.path.isdir(input_path_test) and os.path.isfile(style_path_test):
        style = prep_image_test(misc.imread(style_path_test))
        input_files = os.listdir(input_path_test)
        for input_file in input_files:
            input = prep_image_test(misc.imread(os.path.join(input_path_test, input_file)))
            output = transfer(input, style)
            mon.imwrite('test %s' % input_file[:-4], input)
            mon.imwrite('test %s %s' % (input_file[:-4], style_path_test[:-4]), output)
        mon.imwrite('test %s' % style_path_test[:-4], style)
    else:
        style_files = os.listdir(style_path_test)
        input_files = os.listdir(input_path_test)
        for style_file in style_files:
            style = prep_image_test(misc.imread(os.path.join(style_path_test, style_file)))
            for input_file in input_files:
                input = prep_image_test(misc.imread(os.path.join(input_path_test, input_file)))
                output = transfer(input, style)
                mon.imwrite('test %s' % input_file[:-4], input)
                mon.imwrite('test %s %s' % (input_file[:-4], style_file[:-4]), output)
            mon.imwrite('test %s' % style_file[:-4], style)
    mon.flush()
    print('Testing finished!')
示例#7
0
def train():
    enc = VGG19(input_shape)
    decs = [Decoder(enc, i, name='decoder %d' % i) for i in indices]
    sty_net = StyleTransfer(enc, decs)

    X = T.tensor4('input')
    Y = T.tensor4('style')
    idx = T.scalar('iter', 'int32')
    X_ = nn.placeholder((bs,) + input_shape[1:], name='input_plhd')
    Y_ = nn.placeholder((bs,) + input_shape[1:], name='style_plhd')
    lr_ = nn.placeholder(value=lr, name='lr_plhd')

    nn.set_training_on()
    losses = [dec.cost(X) for dec in decs]
    updates = [nn.adam(loss[0] + weight * loss[1], dec.trainable, lr) for loss, dec in zip(losses, decs)]
    nn.anneal_learning_rate(lr_, idx, 'inverse', decay=decay)
    trains = [nn.function([], [loss[0], loss[1], dec(X, True)], givens={X: X_}, updates=update, name='train decoder')
              for loss, dec, update in zip(losses, decs, updates)]

    nn.set_training_off()
    X_styled = sty_net(X, Y)
    transfer = nn.function([], X_styled, givens={X: X_, Y: Y_}, name='transfer style')

    data_train = DataManager(X_, input_path_train, bs, n_epochs, True, num_val_imgs=num_val_imgs, input_shape=input_shape)
    data_test = DataManagerStyleTransfer((X_, Y_), (input_path_val, style_path_val), bs, 1, input_shape=input_shape)
    mon = nn.Monitor(model_name='WCT', valid_freq=print_freq)

    print('Training...')
    for it in data_train:
        results = [train(it) for train in trains]

        with mon:
            for layer, res in zip(indices, results):
                if np.isnan(res[0] + res[1]) or np.isinf(res[0] + res[1]):
                    raise ValueError('Training failed!')
                mon.plot('pixel loss at layer %d' % layer, res[0])
                mon.plot('feature loss at layer %d' % layer, res[1])

            if it % val_freq == 0:
                mon.imwrite('recon img at layer %d' % layer, res[2])

                for i in data_test:
                    img = transfer()
                    mon.imwrite('stylized image %d' % i, img)
                    mon.imwrite('input %d' % i, X_.get_value())
                    mon.imwrite('style %d' % i, Y_.get_value())

                for idx, dec in zip(indices, decs):
                    mon.dump(nn.utils.shared2numpy(dec.params), 'decoder-%d.npz' % idx, 5)
    mon.flush()
    for idx, dec in zip(indices, decs):
        mon.dump(nn.utils.shared2numpy(dec.params), 'decoder-%d-final.npz' % idx)
    print('Training finished!')
示例#8
0
def main():

    #定义数据加载器
    transform = Transform()

    content_set = CSDataset(transform=transform, root=args.content_dir)
    style_set = CSDataset(transform=transform, root=args.style_dir)

    content_loader = iter(data.DataLoader( content_set, batch_size=args.batch_size,sampler=RecurrentSampler(content_set), \
        num_workers=args.n_threads,pin_memory=True,drop_last=True))

    style_loader = iter(data.DataLoader( style_set, batch_size=args.batch_size,sampler=RecurrentSampler(style_set),\
         num_workers=args.n_threads,pin_memory=True,drop_last=True))

    #初始化模型
    decoder = Decoder()
    decoder = decoder.to(device)
    decoder.zero_grad()

    #训练网络
    train(decoder, content_loader, style_loader)
示例#9
0
def transfer(contentpath,
             stylepath,
             converted,
             pixel=512,
             model_path='static/20200522decoder100000_1.pth'):
    '''一次前传得到风格化图像'''
    mytransfer = test_transform(pixel)

    contentfile = open(str(contentpath), 'rb')
    stylefile = open(str(stylepath), 'rb')

    contentimg = Image.open(contentfile).convert('RGB')
    styleimg = Image.open(stylefile).convert('RGB')

    contentfile.close()
    stylefile.close()

    contentimg = mytransfer(contentimg).unsqueeze(0)
    styleimg = mytransfer(styleimg).unsqueeze(0)

    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = torch.device("cpu")

    decoder = Decoder().to(device).eval()
    decoder.load_state_dict(
        torch.load(model_path, map_location=torch.device('cpu')))
    # decoder = decoder.module
    # decoder.load_state_dict(torch.load(model_path))

    fbnet = FPnet(decoder, True).to(device).eval()
    output = fbnet(contentimg,
                   styleimg,
                   alpha=1.0,
                   lamda=1.0,
                   require_loss=False)

    save_image(output.cpu(), converted)
    contentimg.detach()
    styleimg.detach()
    output.detach()
示例#10
0
def main():
    parser = argparse.ArgumentParser(
        description='chainer implementation of pix2pix')
    parser.add_argument('--batchsize',
                        '-b',
                        type=int,
                        default=1,
                        help='Number of images in each mini-batch')
    parser.add_argument('--epoch',
                        '-e',
                        type=int,
                        default=200,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--gpu',
                        '-g',
                        type=int,
                        default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--out',
                        '-o',
                        default='result',
                        help='Directory to output the result')
    parser.add_argument('--enc-npz', default='', required=True)
    parser.add_argument('--dec-npz', default='', required=True)
    parser.add_argument('input')
    args = parser.parse_args()

    # Set up a neural network to train
    enc = Encoder(in_ch=3)
    dec = Decoder(out_ch=3)

    chainer.serializers.load_npz(args.enc_npz, enc)
    chainer.serializers.load_npz(args.dec_npz, dec)

    # load image and predict

    img = Image.open(args.input).convert('RGB')
    img = np.asarray(img).astype(np.float32)
    img /= 255.0
    img = img.transpose(2, 0, 1).reshape((1, 3, 256, 256))

    x = chainer.Variable(img)
    z = enc(x)
    y = dec(z)  # (batchsize, ch, width, height)

    y = y.data.reshape(3, 256, 256).transpose(1, 2, 0).clip(0, 1)
    import matplotlib.pyplot as plt
    plt.imshow(y)
    plt.show()
    if False:
        import scipy.misc
        scipy.misc.imsave("output.png", y)
示例#11
0
def test_random():
    enc = Encoder((None, ) + input_size, vgg_param_file)
    dec = Decoder(enc.output_shape, dec_param_file)

    X = T.tensor4('input')
    Y = T.tensor4('style')
    weights = [T.vector('weights') for i in range(len(vgg_info))]

    nn.set_training_off()
    X_styled = dec(enc((X, Y), weights))
    test = nn.function([X, Y] + weights, X_styled, name='test generator')

    style_folder = os.listdir(style_img_folder)
    input_folder = os.listdir(input_img_folder)

    time_list = []
    if not os.path.exists('outputs'):
        os.mkdir('outputs')

    for style_file in style_folder:
        sty_img = prep_image_test(
            misc.imread(style_img_folder + '/' + style_file))
        for input_file in input_folder:
            try:
                input_img = prep_image_test(
                    misc.imread(input_img_folder + '/' + input_file))
            except ValueError:
                continue

            for i in range(num_styles):
                start = time.time()
                output = test(input_img, sty_img, *get_weights(vgg_info))
                time_list.append(time.time() - start)
                output = np.transpose(output[0], (1, 2, 0))
                misc.imsave(
                    os.path.join(
                        'outputs', input_file[:-4] + '_' + style_file[:-4] +
                        '_%d.jpg' % i), output)

    print('Took %f s/image' % np.mean(time_list))
    print('Testing finished!')
示例#12
0
def main():
    parser = argparse.ArgumentParser(
        description='chainer implementation of pix2pix')
    parser.add_argument('--batchsize',
                        '-b',
                        type=int,
                        default=1,
                        help='Number of images in each mini-batch')
    parser.add_argument('--epoch',
                        '-e',
                        type=int,
                        default=200,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--gpu',
                        '-g',
                        type=int,
                        default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--out',
                        '-o',
                        default='result',
                        help='Directory to output the result')
    parser.add_argument('--enc-npz', default='', required=True)
    parser.add_argument('--dec-npz', default='', required=True)
    args = parser.parse_args()

    # Set up a neural network to train
    enc = Encoder(in_ch=3)
    dec = Decoder(out_ch=3)

    chainer.serializers.load_npz(args.enc_npz, enc)
    chainer.serializers.load_npz(args.dec_npz, dec)

    # graph
    input = np.zeros((1, 3, 256, 256), dtype=np.float32)
    x = chainer.Variable(input)
    z = enc(x)
    y = dec(z)
    graph = ChainerConverter().convert([x], [y])
    exec_info = generate_descriptor("webassembly", graph)
    exec_info.save(args.out)
示例#13
0
    def __init__(self):
        super(Agent, self).__init__()
        #self.config = config
        self.action = Action()
        self.dataset = Dataset()

        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu")
        self.grad_clip = 5

        embed_dim = 512  # dimension of word embeddings
        decoder_dim = 512  # dimension of decoder RNN
        dropout = 0.5
        #batch_size = 32
        self.n_epochs = 20

        self.encoder = Encoder()
        self.decoder = Decoder(embed_dim=embed_dim,
                               decoder_dim=decoder_dim,
                               vocab_size=self.dataset.get_word_map_len(),
                               dropout=dropout)
示例#14
0
def main():
    parser = argparse.ArgumentParser(
        description='chainer implementation of pix2pix')
    parser.add_argument('--batchsize', '-b', type=int, default=1,
                        help='Number of images in each mini-batch')
    parser.add_argument('--epoch', '-e', type=int, default=200,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--device', '-d', type=str, default='-1',
                        help='Device specifier. Either ChainerX device '
                        'specifier or an integer. If non-negative integer, '
                        'CuPy arrays with specified device id are used. If '
                        'negative integer, NumPy arrays are used')
    parser.add_argument('--dataset', '-i', default='./facade/base',
                        help='Directory of image files.')
    parser.add_argument('--out', '-o', default='result',
                        help='Directory to output the result')
    parser.add_argument('--resume', '-r', type=str,
                        help='Resume the training from snapshot')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed')
    parser.add_argument('--snapshot_interval', type=int, default=1000,
                        help='Interval of snapshot')
    parser.add_argument('--display_interval', type=int, default=100,
                        help='Interval of displaying log to console')
    group = parser.add_argument_group('deprecated arguments')
    group.add_argument('--gpu', '-g', dest='device',
                       type=int, nargs='?', const=0,
                       help='GPU ID (negative value indicates CPU)')
    args = parser.parse_args()

    if chainer.get_dtype() == numpy.float16:
        warnings.warn(
            'This example may cause NaN in FP16 mode.', RuntimeWarning)

    device = chainer.get_device(args.device)
    if device.xp is chainerx:
        sys.stderr.write('This example does not support ChainerX devices.\n')
        sys.exit(1)

    print('Device: {}'.format(device))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))
    print('')

    device.use()

    # Set up a neural network to train
    enc = Encoder(in_ch=12)
    dec = Decoder(out_ch=3)
    dis = Discriminator(in_ch=12, out_ch=3)

    enc.to_device(device)
    dec.to_device(device)
    dis.to_device(device)

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer
    opt_enc = make_optimizer(enc)
    opt_dec = make_optimizer(dec)
    opt_dis = make_optimizer(dis)

    train_d = FacadeDataset(args.dataset, data_range=(1, 300))
    test_d = FacadeDataset(args.dataset, data_range=(300, 379))
    train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
    test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)

    # Set up a trainer
    updater = FacadeUpdater(
        models=(enc, dec, dis),
        iterator={
            'main': train_iter,
            'test': test_iter},
        optimizer={
            'enc': opt_enc, 'dec': opt_dec,
            'dis': opt_dis},
        device=device)
    trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)

    snapshot_interval = (args.snapshot_interval, 'iteration')
    display_interval = (args.display_interval, 'iteration')
    trainer.extend(extensions.snapshot(
        filename='snapshot_iter_{.updater.iteration}.npz'),
        trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.LogReport(trigger=display_interval))
    trainer.extend(extensions.PrintReport([
        'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss',
    ]), trigger=display_interval)
    trainer.extend(extensions.ProgressBar(update_interval=10))
    trainer.extend(
        out_image(
            updater, enc, dec,
            5, 5, args.seed, args.out),
        trigger=snapshot_interval)

    if args.resume is not None:
        # Resume from a snapshot
        chainer.serializers.load_npz(args.resume, trainer)

    # Run the training
    trainer.run()
示例#15
0
文件: train.py 项目: g089vg/python
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed')
    parser.add_argument('--snapshot_interval', type=int, default=1000,
                        help='Interval of snapshot')
    parser.add_argument('--display_interval', type=int, default=100,
                        help='Interval of displaying log to console')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))
    print('')

    # Set up a neural network to train
    enc = Encoder(in_ch=5)
    dec = Decoder(out_ch=3)
    dis = Discriminator(in_ch=5, out_ch=3)
    
    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
        dis.to_gpu()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer
    opt_enc = make_optimizer(enc)
示例#16
0
    return x


def _istft(y):
    _, hop_length, win_length = _stft_parameters()
    return librosa.istft(y, hop_length=hop_length, win_length=win_length)


def save_wav(wav, path):
    wav *= 1 / max(0.01, np.max(np.abs(wav)))
    librosa.output.write_wav(path, wav, hparams.sample_rate)


# Set up a neural network
enc = Encoder(in_ch=2)
dec = Decoder(out_ch=2)

# load terained models
chainer.serializers.load_npz(MODEL_DIR + "enc_iter_%d.npz" % ITER, enc)
chainer.serializers.load_npz(MODEL_DIR + "dec_iter_%d.npz" % ITER, dec)

#img = Image.open('../input/png/%d_ritsuko.png'%MUSIC_IDX)
#img = np.asarray(img).astype("f")/128.0-1.0
img = sp2.astype("f") / 128.0 - 1.0
h, w, c = img.shape
print(img.shape)
img = img.transpose(2, 0, 1)

img = img[:2, :, :]
img = img.reshape(1, 2, h, w)
print(img.shape)
示例#17
0
def main():
    parser = argparse.ArgumentParser(
        description='chainer creating pictures of seat')
    parser.add_argument('--gpu',
                        '-g',
                        type=int,
                        default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--data',
                        '-d',
                        default='./detect',
                        help='Directory of image files.')
    parser.add_argument('--mask',
                        '-ma',
                        default='./mask',
                        help='Directory of image files.')
    parser.add_argument('--out',
                        '-o',
                        default='./result',
                        help='Directory to output the result')
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    parser.add_argument('--model',
                        '-m',
                        default='snapshot_iter_83000.npz',
                        help='Loading model')
    parser.add_argument('--batchsize',
                        '-b',
                        default=16,
                        type=int,
                        help='The same value as that of trainer')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('')

    # Set up a neural network to train
    enc = Encoder(in_ch=4)
    dec = Decoder(out_ch=3)

    if args.gpu >= 0:
        cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()

    xp = cuda.cupy

    # Setup an model
    model = args.model
    print('---Loading model---')
    serializers.load_npz(model, enc, path='updater/model:enc/')
    serializers.load_npz(model, dec, path='updater/model:dec/')

    # Trun off a trainer
    chainer.config.train = False
    chainer.config.debug = False
    chainer.config.enable_backprop = False
    print('Setuped model!')

    # Load datasets
    print('---Loading datasets---')
    data_path = args.data
    data_sets = glob.glob(data_path + '/*.jpg')
    data_mask = glob.glob(args.mask + '/*.jpg')
    dataset = []
    names = []
    for data in data_sets:
        d_name = os.path.basename(data)
        d_name = d_name[:-4]
        img = Image.open(data)
        img = xp.asarray(img).transpose(2, 0, 1)
        for _ in range(10):
            mask = random.choice(data_mask)
            mask = Image.open(mask)
            mask = xp.asarray(mask)
            mask = mask[xp.newaxis, :, :]
            img_ = img + mask
            img_ = xp.asarray(img_).astype('f') / 128.0 - 1.0
            mask = xp.asarray(mask).astype('f') / 128.0 - 1.0
            img_ = xp.concatenate([img_, mask], axis=0)
            dataset.append(img_)
            f_name = d_name + '_' + str(_)
            names.append(f_name)
    print('Setuped datasets!')

    # Create picture
    print('---Creating---')
    in_ch = 4
    in_h = 256
    in_w = 256
    out_put = 0

    batch_size = args.batchsize
    out_dir = args.out

    if not (os.path.exists(out_dir)):
        os.mkdir(out_dir)

    _ = 0
    for name, data in zip(names, dataset):
        X_in = xp.zeros((batch_size, in_ch, in_h, in_w)).astype("f")
        for i in range(batch_size):
            X_in[i, :] = xp.asarray(data)
        X_in = Variable(X_in)

        z = enc(X_in)
        X_out = dec(z)
        out_put = xp.asarray(X_out.data)
        out_put = out_put[0]
        out_put += 1.0
        out_put *= 128.0
        xp.save(out_dir + '/' + name, out_put)
        _ += 1
        print('created {} / {}'.format(_, len(dataset)))

    print('Finished all process!')
    print('Numpy shape : ', out_put.shape)
    print('Number of Numpy file : ', len(dataset))
示例#18
0
def main():
    parser = argparse.ArgumentParser(
        description="chainer implementation of Unet")
    parser.add_argument("--batchsize",
                        "-b",
                        type=int,
                        default=1,
                        help="Number of images in each mini-batch")
    parser.add_argument("--epoch", "-e", type=int, default=200, help="epoch")
    parser.add_argument("--gpu", "-g", type=int, default=-1, help="GPU ID")
    parser.add_argument("--dataset",
                        "-i",
                        default="./train/",
                        help="Directory of image files")
    parser.add_argument("--out",
                        "-o",
                        default="result/",
                        help="Directory to output the result")
    parser.add_argument('--resume',
                        '-r',
                        default='',
                        help='Resume the training from snapshot')
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    parser.add_argument('--snapshot_interval',
                        type=int,
                        default=1000,
                        help='Interval of snapshot')
    parser.add_argument('--display_interval',
                        type=int,
                        default=100,
                        help='Interval of displaying log to console')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))
    print('')

    # Only Encoder Decoder with Unet
    enc = Encoder(in_ch=3)  # in_ch => 3(YCbCr)
    dec = Decoder(out_ch=3)  # out_ch => 3(DCT)

    # GPU set up
    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()
        enc.to_gpu()
        dec.to_gpu()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer

    opt_enc = make_optimizer(enc)
    opt_dec = make_optimizer(dec)

    train_d = ImgDCTDataset(args.dataset, data_range=(0, 1000))
    test_d = ImgDCTDataset(args.dataset, data_range=(1000, 2000))
    train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
    test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)

    updater = FacadeUpdater(models=(enc, dec),
                            iterator={
                                'main': train_iter,
                                'test': test_iter
                            },
                            optimizer={
                                'enc': opt_enc,
                                'dec': opt_dec
                            },
                            device=args.gpu)
    trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)

    snapshot_interval = (args.snapshot_interval, 'iteration')
    display_interval = (args.display_interval, 'iteration')
    trainer.extend(
        extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'),
        trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        enc, 'enc_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dec, 'dec_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    # trainer.extend(extensions.snapshot_object(
    #     dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.LogReport(trigger=display_interval))
    trainer.extend(extensions.PrintReport([
        'epoch',
        'iteration',
        'enc/loss',
        'dec/loss',
    ]),
                   trigger=display_interval)
    trainer.extend(extensions.ProgressBar(update_interval=10))

    if args.resume:
        # Resume from a snapshot
        chainer.serializers.load_npz(args.resume, trainer)

    # Run the training
    trainer.run()
示例#19
0
def main():
    parser = argparse.ArgumentParser(
        description='chainer implementation of pix2pix')
    parser.add_argument('--batchsize',
                        '-b',
                        type=int,
                        default=1,
                        help='Number of images in each mini-batch')
    parser.add_argument('--epoch',
                        '-e',
                        type=int,
                        default=200,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--gpu',
                        '-g',
                        type=int,
                        default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--dataset',
                        '-i',
                        default='./facade/base',
                        help='Directory of image files.')
    parser.add_argument('--out',
                        '-o',
                        default='result',
                        help='Directory to output the result')
    parser.add_argument('--resume',
                        '-r',
                        default='',
                        help='Resume the training from snapshot')
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    parser.add_argument('--snapshot_interval',
                        type=int,
                        default=1000,
                        help='Interval of snapshot')
    parser.add_argument('--display_interval',
                        type=int,
                        default=100,
                        help='Interval of displaying log to console')
    parser.add_argument(
        '--n_processes',
        type=int,
        default=None,
        help='processes of chainer.iterators.MultiprocessIterator')
    parser.add_argument(
        '--shared_mem',
        type=int,
        default=None,
        help=
        'shared memory per data, for chainer.iterators.MultiprocessIterator. None means auto ajust.'
    )
    parser.add_argument('--audio_dataset_second',
                        type=int,
                        default=None,
                        help='time length(second) of train audio data .')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))
    print('')

    args.out = os.path.join(args.out, datetime.now().strftime("%Y%m%d_%H%M%S"))
    util.audio_dataset_second = args.audio_dataset_second
    if args.batchsize > 1:
        assert util.audio_dataset_second != None, "when minibatch training (e.g. --batchsize > 1), --audio_dataset_second option is required."

    # Set up a neural network to train
    enc = Encoder(in_ch=2)
    dec = Decoder(out_ch=2)
    dis = Discriminator(in_ch=2, out_ch=2)

    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
        dis.to_gpu()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer

    opt_enc = make_optimizer(enc)
    opt_dec = make_optimizer(dec)
    opt_dis = make_optimizer(dis)

    train_d = Vp2pDataset(args.dataset + "/train")
    test_d = Vp2pDataset(args.dataset + "/test")
    train_iter = chainer.iterators.MultiprocessIterator(
        train_d,
        args.batchsize,
        n_processes=args.n_processes,
        shared_mem=args.shared_mem)
    test_iter = chainer.iterators.MultiprocessIterator(
        test_d,
        args.batchsize,
        n_processes=args.n_processes,
        shared_mem=args.shared_mem)
    # train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize)
    # test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize)

    # Set up a trainer
    updater = VoiceP2PUpdater(models=(enc, dec, dis),
                              iterator={
                                  'main': train_iter,
                                  'test': test_iter
                              },
                              optimizer={
                                  'enc': opt_enc,
                                  'dec': opt_dec,
                                  'dis': opt_dis
                              },
                              device=args.gpu)
    trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)

    snapshot_interval = (args.snapshot_interval, 'iteration')
    display_interval = (args.display_interval, 'iteration')
    trainer.extend(
        extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'),
        trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        enc, 'enc_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dec, 'dec_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dis, 'dis_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.LogReport(trigger=display_interval))
    trainer.extend(extensions.PrintReport([
        'epoch',
        'iteration',
        'enc/loss',
        'dec/loss',
        'dis/loss',
    ]),
                   trigger=display_interval)
    trainer.extend(extensions.ProgressBar(update_interval=10))
    trainer.extend(out_image(updater, enc, dec, 5, 5, args.seed, args.out),
                   trigger=snapshot_interval)

    if args.resume:
        # Resume from a snapshot
        chainer.serializers.load_npz(args.resume, trainer)

    # Run the training
    trainer.run()
示例#20
0
def main():
    parser = argparse.ArgumentParser(
        description='chainer implementation of pix2pix')

    parser.add_argument('--gpu',
                        '-g',
                        type=int,
                        default=-1,
                        help='GPU ID (negative value indicates CPU)')

    parser.add_argument('--seed', type=int, default=0, help='Random seed')

    parser.add_argument('--model', '-m', default='', help='model snapshot')

    parser.add_argument('--input',
                        '-i',
                        default='../images/generate/sample.jpg',
                        help='input jpg')

    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))

    # Set up a neural network to train
    enc = Encoder(in_ch=3)
    dec = Decoder(out_ch=3)
    dis = Discriminator(in_ch=3, out_ch=3)

    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
        dis.to_gpu()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer

    opt_enc = make_optimizer(enc)
    opt_dec = make_optimizer(dec)
    opt_dis = make_optimizer(dis)

    if os.path.exists('generate_tmp'):
        shutil.rmtree('generate_tmp')

    os.mkdir('generate_tmp')
    os.mkdir('generate_tmp/base')
    os.mkdir('generate_tmp/label')
    shutil.copyfile(args.input, 'generate_tmp/base/tmp.jpg')
    shutil.copyfile(args.input, 'generate_tmp/label/tmp.jpg')
    test_d = FacadeDataset('generate_tmp/')
    test_iter = chainer.iterators.SerialIterator(test_d, 1)

    # Set up a trainer
    updater = FacadeUpdater(models=(enc, dec, dis),
                            iterator={},
                            optimizer={
                                'enc': opt_enc,
                                'dec': opt_dec,
                                'dis': opt_dis
                            },
                            device=args.gpu)
    trainer = training.Trainer(updater, (200, 'epoch'),
                               out='../results/generate/')
    chainer.serializers.load_npz(args.model, trainer)

    out_image(updater, enc, dec, 1, 1, args.seed, '../results/generate/', True,
              test_iter)(trainer)
示例#21
0
def main():
    args = arguments()
    print(args)

    if args.imgtype == "dcm":
        from dataset_dicom import Dataset as Dataset
    else:
        from dataset_jpg import DatasetOutMem as Dataset

    # CUDA
    if not chainer.cuda.available:
        print("CUDA required")
        exit()
    if len(args.gpu) == 1 and args.gpu[0] >= 0:
        chainer.cuda.get_device_from_id(args.gpu[0]).use()

    # Enable autotuner of cuDNN
    chainer.config.autotune = True
    chainer.config.dtype = dtypes[args.dtype]
    chainer.print_runtime_info()
    # Turn off type check
    #    chainer.config.type_check = False
    #    print('Chainer version: ', chainer.__version__)
    #    print('GPU availability:', chainer.cuda.available)
    #    print('cuDNN availablility:', chainer.cuda.cudnn_enabled)

    ## dataset iterator
    print("Setting up data iterators...")
    train_A_dataset = Dataset(path=os.path.join(args.root, 'trainA'),
                              args=args,
                              base=args.HU_baseA,
                              rang=args.HU_rangeA,
                              random=args.random_translate)
    train_B_dataset = Dataset(path=os.path.join(args.root, 'trainB'),
                              args=args,
                              base=args.HU_baseB,
                              rang=args.HU_rangeB,
                              random=args.random_translate)
    test_A_dataset = Dataset(path=os.path.join(args.root, 'testA'),
                             args=args,
                             base=args.HU_baseA,
                             rang=args.HU_rangeA,
                             random=0)
    test_B_dataset = Dataset(path=os.path.join(args.root, 'testB'),
                             args=args,
                             base=args.HU_baseB,
                             rang=args.HU_rangeB,
                             random=0)

    args.ch = train_A_dataset.ch
    args.out_ch = train_B_dataset.ch
    print("channels in A {}, channels in B {}".format(args.ch, args.out_ch))

    #    test_A_iter = chainer.iterators.SerialIterator(test_A_dataset, args.nvis_A, shuffle=False)
    #    test_B_iter = chainer.iterators.SerialIterator(test_B_dataset, args.nvis_B, shuffle=False)
    test_A_iter = chainer.iterators.MultithreadIterator(test_A_dataset,
                                                        args.nvis_A,
                                                        shuffle=False,
                                                        n_threads=3)
    test_B_iter = chainer.iterators.MultithreadIterator(test_B_dataset,
                                                        args.nvis_B,
                                                        shuffle=False,
                                                        n_threads=3)
    train_A_iter = chainer.iterators.MultithreadIterator(train_A_dataset,
                                                         args.batch_size,
                                                         n_threads=3)
    train_B_iter = chainer.iterators.MultithreadIterator(train_B_dataset,
                                                         args.batch_size,
                                                         n_threads=3)

    # setup models
    enc_x = Encoder(args)
    enc_y = enc_x if args.single_encoder else Encoder(args)
    dec_x = Decoder(args)
    dec_y = Decoder(args)
    dis_x = Discriminator(args)
    dis_y = Discriminator(args)
    dis_z = Discriminator(
        args) if args.lambda_dis_z > 0 else chainer.links.Linear(1, 1)
    models = {
        'enc_x': enc_x,
        'dec_x': dec_x,
        'enc_y': enc_y,
        'dec_y': dec_y,
        'dis_x': dis_x,
        'dis_y': dis_y,
        'dis_z': dis_z
    }

    ## load learnt models
    if args.load_models:
        for e in models:
            m = args.load_models.replace('enc_x', e)
            try:
                serializers.load_npz(m, models[e])
                print('model loaded: {}'.format(m))
            except:
                print("couldn't load {}".format(m))
                pass

    # select GPU
    if len(args.gpu) == 1:
        for e in models:
            models[e].to_gpu()
        print('using gpu {}, cuDNN {}'.format(args.gpu,
                                              chainer.cuda.cudnn_enabled))
    else:
        print("mandatory GPU use: currently only a single GPU can be used")
        exit()

    # Setup optimisers
    def make_optimizer(model, lr, opttype='Adam'):
        #        eps = 1e-5 if args.dtype==np.float16 else 1e-8
        optimizer = optim[opttype](lr)
        #from profiled_optimizer import create_marked_profile_optimizer
        #        optimizer = create_marked_profile_optimizer(optim[opttype](lr), sync=True, sync_level=2)
        optimizer.setup(model)
        if args.weight_decay > 0:
            if opttype in ['Adam', 'AdaBound', 'Eve']:
                optimizer.weight_decay_rate = args.weight_decay
            else:
                if args.weight_decay_norm == 'l2':
                    optimizer.add_hook(
                        chainer.optimizer.WeightDecay(args.weight_decay))
                else:
                    optimizer.add_hook(
                        chainer.optimizer_hooks.Lasso(args.weight_decay))
        return optimizer

    opt_enc_x = make_optimizer(enc_x, args.learning_rate_g, args.optimizer)
    opt_dec_x = make_optimizer(dec_x, args.learning_rate_g, args.optimizer)
    opt_enc_y = make_optimizer(enc_y, args.learning_rate_g, args.optimizer)
    opt_dec_y = make_optimizer(dec_y, args.learning_rate_g, args.optimizer)
    opt_x = make_optimizer(dis_x, args.learning_rate_d, args.optimizer)
    opt_y = make_optimizer(dis_y, args.learning_rate_d, args.optimizer)
    opt_z = make_optimizer(dis_z, args.learning_rate_d, args.optimizer)
    optimizers = {
        'opt_enc_x': opt_enc_x,
        'opt_dec_x': opt_dec_x,
        'opt_enc_y': opt_enc_y,
        'opt_dec_y': opt_dec_y,
        'opt_x': opt_x,
        'opt_y': opt_y,
        'opt_z': opt_z
    }
    if args.load_optimizer:
        for e in optimizers:
            try:
                m = args.load_models.replace('enc_x', e)
                serializers.load_npz(m, optimizers[e])
                print('optimiser loaded: {}'.format(m))
            except:
                print("couldn't load {}".format(m))
                pass

    # Set up an updater: TODO: multi gpu updater
    print("Preparing updater...")
    updater = Updater(models=(enc_x, dec_x, enc_y, dec_y, dis_x, dis_y, dis_z),
                      iterator={
                          'main': train_A_iter,
                          'train_B': train_B_iter,
                      },
                      optimizer=optimizers,
                      converter=convert.ConcatWithAsyncTransfer(),
                      device=args.gpu[0],
                      params={'args': args})

    if args.snapinterval < 0:
        args.snapinterval = args.lrdecay_start + args.lrdecay_period
    log_interval = (200, 'iteration')
    model_save_interval = (args.snapinterval, 'epoch')
    plot_interval = (500, 'iteration')

    # Set up a trainer
    print("Preparing trainer...")
    if args.iteration:
        stop_trigger = (args.iteration, 'iteration')
    else:
        stop_trigger = (args.lrdecay_start + args.lrdecay_period, 'epoch')
    trainer = training.Trainer(updater, stop_trigger, out=args.out)
    for e in models:
        trainer.extend(extensions.snapshot_object(models[e],
                                                  e + '{.updater.epoch}.npz'),
                       trigger=model_save_interval)


#        trainer.extend(extensions.ParameterStatistics(models[e]))   ## very slow
    for e in optimizers:
        trainer.extend(extensions.snapshot_object(optimizers[e],
                                                  e + '{.updater.epoch}.npz'),
                       trigger=model_save_interval)

    log_keys = ['epoch', 'iteration', 'lr']
    log_keys_cycle = [
        'opt_enc_x/loss_cycle', 'opt_enc_y/loss_cycle', 'opt_dec_x/loss_cycle',
        'opt_dec_y/loss_cycle', 'myval/cycle_x_l1', 'myval/cycle_y_l1'
    ]
    log_keys_adv = [
        'opt_enc_y/loss_adv', 'opt_dec_y/loss_adv', 'opt_enc_x/loss_adv',
        'opt_dec_x/loss_adv'
    ]
    log_keys_d = []
    if args.lambda_reg > 0:
        log_keys.extend(['opt_enc_x/loss_reg', 'opt_enc_y/loss_reg'])
    if args.lambda_tv > 0:
        log_keys.extend(['opt_dec_y/loss_tv'])
    if args.lambda_air > 0:
        log_keys.extend(['opt_dec_x/loss_air', 'opt_dec_y/loss_air'])
    if args.lambda_grad > 0:
        log_keys.extend(['opt_dec_x/loss_grad', 'opt_dec_y/loss_grad'])
    if args.lambda_identity_x > 0:  # perceptual
        log_keys.extend(['opt_enc_x/loss_id', 'opt_enc_y/loss_id'])
    if args.lambda_domain > 0:
        log_keys_cycle.extend(['opt_dec_x/loss_dom', 'opt_dec_y/loss_dom'])
    if args.dis_reg_weighting > 0:
        log_keys_d.extend(
            ['opt_x/loss_reg', 'opt_y/loss_reg', 'opt_z/loss_reg'])
    if args.dis_wgan:
        log_keys_d.extend([
            'opt_x/loss_dis', 'opt_x/loss_gp', 'opt_y/loss_dis',
            'opt_y/loss_gp'
        ])
        if args.lambda_dis_z > 0:
            log_keys_d.extend(['opt_z/loss_dis', 'opt_z/loss_gp'])
    else:
        log_keys_d.extend([
            'opt_x/loss_real', 'opt_x/loss_fake', 'opt_y/loss_real',
            'opt_y/loss_fake'
        ])
        if args.lambda_dis_z > 0:
            log_keys_d.extend(['opt_z/loss_x', 'opt_z/loss_y'])

    log_keys_all = log_keys + log_keys_d + log_keys_adv + log_keys_cycle
    trainer.extend(
        extensions.LogReport(keys=log_keys_all, trigger=log_interval))
    trainer.extend(extensions.PrintReport(log_keys_all), trigger=log_interval)
    trainer.extend(extensions.ProgressBar(update_interval=20))
    trainer.extend(extensions.observe_lr(optimizer_name='opt_enc_x'),
                   trigger=log_interval)
    # learning rate scheduling
    decay_start_iter = len(train_A_dataset) * args.lrdecay_start
    decay_end_iter = len(train_A_dataset) * (args.lrdecay_start +
                                             args.lrdecay_period)
    for e in [opt_enc_x, opt_enc_y, opt_dec_x, opt_dec_y]:
        trainer.extend(
            extensions.LinearShift('alpha', (args.learning_rate_g, 0),
                                   (decay_start_iter, decay_end_iter),
                                   optimizer=e))
    for e in [opt_x, opt_y, opt_z]:
        trainer.extend(
            extensions.LinearShift('alpha', (args.learning_rate_d, 0),
                                   (decay_start_iter, decay_end_iter),
                                   optimizer=e))
    ## dump graph
    if args.lambda_Az > 0:
        trainer.extend(
            extensions.dump_graph('opt_enc_x/loss_cycle', out_name='gen.dot'))
    if args.lambda_dis_x > 0:
        if args.dis_wgan:
            trainer.extend(
                extensions.dump_graph('opt_x/loss_dis', out_name='dis.dot'))
        else:
            trainer.extend(
                extensions.dump_graph('opt_x/loss_fake', out_name='dis.dot'))

    # ChainerUI
    trainer.extend(CommandsExtension())

    if extensions.PlotReport.available():
        trainer.extend(
            extensions.PlotReport(log_keys[3:],
                                  'iteration',
                                  trigger=plot_interval,
                                  file_name='loss.png',
                                  postprocess=plot_log))
        trainer.extend(
            extensions.PlotReport(log_keys_d,
                                  'iteration',
                                  trigger=plot_interval,
                                  file_name='loss_d.png'))
        trainer.extend(
            extensions.PlotReport(log_keys_adv,
                                  'iteration',
                                  trigger=plot_interval,
                                  file_name='loss_adv.png'))
        trainer.extend(
            extensions.PlotReport(log_keys_cycle,
                                  'iteration',
                                  trigger=plot_interval,
                                  file_name='loss_cyc.png',
                                  postprocess=plot_log))

    ## visualisation
    vis_folder = os.path.join(args.out, "vis")
    os.makedirs(vis_folder, exist_ok=True)
    if not args.vis_freq:
        args.vis_freq = len(train_A_dataset) // 2
    s = [k for k in range(args.num_slices)
         ] if args.num_slices > 0 and args.imgtype == "dcm" else None
    trainer.extend(VisEvaluator({
        "testA": test_A_iter,
        "testB": test_B_iter
    }, {
        "enc_x": enc_x,
        "enc_y": enc_y,
        "dec_x": dec_x,
        "dec_y": dec_y
    },
                                params={
                                    'vis_out': vis_folder,
                                    'slice': s,
                                    'args': args
                                },
                                device=args.gpu[0]),
                   trigger=(args.vis_freq, 'iteration'))

    ## output filenames of training dataset
    with open(os.path.join(args.out, 'trainA.txt'), 'w') as output:
        for f in train_A_dataset.names:
            output.writelines("\n".join(f))
            output.writelines("\n")
    with open(os.path.join(args.out, 'trainB.txt'), 'w') as output:
        for f in train_B_dataset.names:
            output.writelines("\n".join(f))
            output.writelines("\n")

    # archive the scripts
    rundir = os.path.dirname(os.path.realpath(__file__))
    import zipfile
    with zipfile.ZipFile(os.path.join(args.out, 'script.zip'),
                         'w',
                         compression=zipfile.ZIP_DEFLATED) as new_zip:
        for f in [
                'train.py', 'net.py', 'updater.py', 'consts.py', 'losses.py',
                'arguments.py', 'convert.py'
        ]:
            new_zip.write(os.path.join(rundir, f), arcname=f)

    # Run the training
    print("\nresults are saved under: ", args.out)
    save_args(args, args.out)
    with open(os.path.join(args.out, "args.txt"), 'w') as fh:
        fh.write(" ".join(sys.argv))
    trainer.run()
示例#22
0
def main():
    parser = argparse.ArgumentParser(
        description="chainer implementation of pix2pix")
    parser.add_argument("--batchsize",
                        "-b",
                        type=int,
                        default=1,
                        help="Number of images in each mini-batch")
    parser.add_argument("--epoch",
                        "-e",
                        type=int,
                        default=40000,
                        help="Number of sweeps over the dataset to train")
    parser.add_argument("--gpu",
                        "-g",
                        type=int,
                        default=-1,
                        help="GPU ID (negative value indicates CPU)")
    parser.add_argument("--dataset",
                        "-i",
                        default="./input/png/",
                        help="Directory of image files.")
    parser.add_argument("--out",
                        "-o",
                        default="D:/output/imasUtaConverter/",
                        help="Directory to output the result")
    parser.add_argument("--resume",
                        "-r",
                        default="",
                        help="Resume the training from snapshot")
    parser.add_argument("--seed", type=int, default=0, help="Random seed")
    parser.add_argument("--snapshot_interval",
                        type=int,
                        default=10000,
                        help="Interval of snapshot")
    parser.add_argument("--display_interval",
                        type=int,
                        default=20,
                        help="Interval of displaying log to console")
    args = parser.parse_args()

    print("GPU: {}".format(args.gpu))
    print("# Minibatch-size: {}".format(args.batchsize))
    print("# epoch: {}".format(args.epoch))
    print("")

    # Set up a neural network to train
    enc = Encoder(in_ch=2)
    dec = Decoder(out_ch=2)
    dis = Discriminator(in_ch=2, out_ch=2)

    if args.gpu >= 0:
        chainer.backends.cuda.get_device_from_id(
            args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
        dis.to_gpu()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), "hook_dec")
        return optimizer

    opt_enc = make_optimizer(enc)
    opt_dec = make_optimizer(dec)
    opt_dis = make_optimizer(dis)

    train_d = FacadeDataset(args.dataset, data_range=(0, 38))
    test_d = FacadeDataset(args.dataset, data_range=(38, 40))
    #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4)
    #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4)
    train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
    test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)

    # Set up a trainer
    updater = FacadeUpdater(models=(enc, dec, dis),
                            iterator={
                                "main": train_iter,
                                "test": test_iter
                            },
                            optimizer={
                                "enc": opt_enc,
                                "dec": opt_dec,
                                "dis": opt_dis
                            },
                            device=args.gpu)
    trainer = training.Trainer(updater, (args.epoch, "epoch"), out=args.out)

    snapshot_interval = (args.snapshot_interval, "iteration")
    display_interval = (args.display_interval, "iteration")
    trainer.extend(
        extensions.snapshot(filename="snapshot_iter_{.updater.iteration}.npz"),
        trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        enc, "enc_iter_{.updater.iteration}.npz"),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dec, "dec_iter_{.updater.iteration}.npz"),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dis, "dis_iter_{.updater.iteration}.npz"),
                   trigger=snapshot_interval)
    trainer.extend(extensions.LogReport(trigger=display_interval))
    trainer.extend(extensions.PrintReport([
        "epoch",
        "iteration",
        "enc/loss",
        "dec/loss",
        "dis/loss",
    ]),
                   trigger=display_interval)
    #trainer.extend(extensions.PlotReport(["enc/loss", "dis/loss"], x_key="epoch", file_name="loss.png"))
    trainer.extend(extensions.ProgressBar(update_interval=20))
    trainer.extend(out_image(updater, enc, dec, 1, 10, args.seed, args.out),
                   trigger=snapshot_interval)

    if args.resume:
        # Resume from a snapshot
        chainer.serializers.load_npz(args.resume, trainer)

    # Run the training
    trainer.run()
示例#23
0
def main():
    parser = argparse.ArgumentParser(
        description='chainer implementation of pix2pix')
    parser.add_argument('--batchsize',
                        '-b',
                        type=int,
                        default=1,
                        help='Number of images in each mini-batch')
    parser.add_argument('--epoch',
                        '-e',
                        type=int,
                        default=500,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--gpu',
                        '-g',
                        type=int,
                        default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--out',
                        '-o',
                        default='result_dehighlight',
                        help='Directory to output the result')
    parser.add_argument('--resume',
                        '-r',
                        default='',
                        help='Resume the training from snapshot')
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    parser.add_argument('--snapshot_interval',
                        type=int,
                        default=10000,
                        help='Interval of snapshot')
    parser.add_argument('--display_interval',
                        type=int,
                        default=10000,
                        help='Interval of displaying log to console')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))
    print('')

    # Set up a neural network to train
    enc = Encoder(in_ch=3)
    dec = Decoder(out_ch=1)
    dis = Discriminator(in_ch=3, out_ch=1)
    gen = Generator(in_ch=3, out_ch=1)
    serializers.load_npz("depro.npz", depro)
    gen.encoder = enc
    gen.decoder = dec

    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
        dis.to_gpu()
        gen.to_gpu()

    depro.disable_update()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer

    opt_enc = make_optimizer(enc)
    opt_dec = make_optimizer(dec)
    opt_dis = make_optimizer(dis)

    #will fix
    train_d = NyuDataset("E:/nyu_depth_v2_labeled.mat",
                         startnum=0,
                         endnum=1000)
    test_d = NyuDataset("E:/nyu_depth_v2_labeled.mat",
                        startnum=1000,
                        endnum=1449)
    #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4)
    #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4)
    train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
    test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)
    test_iter2 = chainer.iterators.SerialIterator(test_d,
                                                  args.batchsize,
                                                  repeat=False,
                                                  shuffle=False)

    # Set up a trainer
    updater = PicUpdater(models=(enc, dec, dis),
                         iterator={
                             'main': train_iter,
                             'test': test_iter
                         },
                         optimizer={
                             'enc': opt_enc,
                             'dec': opt_dec,
                             'dis': opt_dis
                         },
                         device=args.gpu)
    trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)

    snapshot_interval = (args.snapshot_interval, 'iteration')
    display_interval = (args.display_interval, 'iteration')
    trainer.extend(
        extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'),
        trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        enc, 'enc_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dec, 'dec_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dis, 'dis_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.LogReport(trigger=display_interval))
    trainer.extend(extensions.PrintReport([
        'epoch', 'iteration', 'enc/loss_enc', 'dec/loss_dec', 'dis/loss_dis',
        "validation/main/loss"
    ]),
                   trigger=display_interval)
    trainer.extend(extensions.Evaluator(test_iter2, gen, device=args.gpu))
    trainer.extend(
        extensions.PlotReport(['dec/loss_dec'],
                              x_key='iteration',
                              file_name='dec_loss.png',
                              trigger=display_interval))
    trainer.extend(
        extensions.PlotReport(['dis/loss_dis'],
                              x_key='iteration',
                              file_name='dis_loss.png',
                              trigger=display_interval))
    trainer.extend(
        extensions.PlotReport(["validation/main/loss"],
                              x_key='iteration',
                              file_name='gen_loss.png',
                              trigger=display_interval))
    trainer.extend(extensions.ProgressBar(update_interval=10))
    trainer.extend(out_image(updater, depro, enc, dec, 3, 3, args.seed,
                             args.out),
                   trigger=snapshot_interval)

    if args.resume:
        # Resume from a snapshot
        chainer.serializers.load_npz(args.resume, trainer)

    # Run the training
    trainer.run()
def main():
    parser = argparse.ArgumentParser(description='chainer implementation of pix2pix')
    parser.add_argument('--batchsize', '-b', type=int, default=1,
                        help='Number of images in each mini-batch')
    parser.add_argument('--epoch', '-e', type=int, default=200,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--gpu', '-g', type=int, default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--dataset', '-i', default='./facade/base',
                        help='Directory of image files.')
    parser.add_argument('--out', '-o', default='result',
                        help='Directory to output the result')
    parser.add_argument('--resume', '-r', default='',
                        help='Resume the training from snapshot')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed')
    parser.add_argument('--snapshot_interval', type=int, default=1000,
                        help='Interval of snapshot')
    parser.add_argument('--display_interval', type=int, default=100,
                        help='Interval of displaying log to console')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))
    print('')

    # Set up a neural network to train
    enc = Encoder(in_ch=12)
    dec = Decoder(out_ch=3)
    dis = Discriminator(in_ch=12, out_ch=3)
    
    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
        dis.to_gpu()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer
    opt_enc = make_optimizer(enc)
    opt_dec = make_optimizer(dec)
    opt_dis = make_optimizer(dis)

    train_d = FacadeDataset(args.dataset, data_range=(1,300))
    test_d = FacadeDataset(args.dataset, data_range=(300,379))
    #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4)
    #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4)
    train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
    test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)

    # Set up a trainer
    updater = FacadeUpdater(
        models=(enc, dec, dis),
        iterator={
            'main': train_iter,
            'test': test_iter},
        optimizer={
            'enc': opt_enc, 'dec': opt_dec, 
            'dis': opt_dis},
        device=args.gpu)
    trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)

    snapshot_interval = (args.snapshot_interval, 'iteration')
    display_interval = (args.display_interval, 'iteration')
    trainer.extend(extensions.snapshot(
        filename='snapshot_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.LogReport(trigger=display_interval))
    trainer.extend(extensions.PrintReport([
        'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss',
    ]), trigger=display_interval)
    trainer.extend(extensions.ProgressBar(update_interval=10))
    trainer.extend(
        out_image(
            updater, enc, dec,
            5, 5, args.seed, args.out),
        trigger=snapshot_interval)

    if args.resume:
        # Resume from a snapshot
        chainer.serializers.load_npz(args.resume, trainer)

    # Run the training
    trainer.run()
示例#25
0
def main():
    parser = argparse.ArgumentParser(description='chainer implementation of pix2pix')
    parser.add_argument('--gpu', '-g', type=int, default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--img', '-i', help='Input image')

    parser.add_argument('--out', '-o', default='result_dehighlight',
                        help='Directory to output the result')
    args = parser.parse_args()

    ENC_W = os.path.join(args.out, "enc_iter_2500000.npz")
    #DEC_W = "trained_model/dec_iter_176000.npz"
    # to avoid GitHub 100M limit, one .npz files are divided into two zip files.
    DEC_Ws = [os.path.join(args.out, "dec_iter_2500000.npz")]

    #shutil.copy("net.py", args.out)

    # Set up a neural network to train
    enc = Encoder(in_ch=3)
    dec = Decoder(out_ch=3)

    chainer.serializers.load_npz(ENC_W, enc)
    # to avoid GitHub 100M limit, 1 .npz file is devided into 2 files
    for npzfile in DEC_Ws:
        with np.load(npzfile) as f:
            d = NpzDeserializer(f, strict=False)
            d.load(dec)

    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()

    inimg = loadimg(args.img)
    ch, h, w = inimg.shape
    # add paddings so that input array has the size of mutiples of 256.
    in_ary = np.zeros((ch,math.ceil(h/256)*256, math.ceil(w/256)*256), dtype="f")
    in_ary[:,0:h,0:w] = inimg
    x_in = in_ary[np.newaxis,:] # to fit into the minibatch shape
    print(x_in.shape)
    # x_in as an input image
    x_in = chainer.Variable(x_in)
    if args.gpu >= 0:
        x_in.to_gpu()

    st = time.time()
    for i in tqdm(range(10)):
        z = enc(x_in)
        x_out = dec(z)
    ts = (time.time() - st)/10
    print("mean estimation time:{:.2f}".format(ts))
    with open(os.path.join(args.out, "time.txt"), "a") as f:
        f.write("gpu:{}, time:{:.4f}, FPS:{:.4f}\n".format(args.gpu, ts, 1/ts))

    if args.gpu >= 0:
        out_ary = x_out.data.get()[0]
    else:
        out_ary = x_out.data[0]
    #img_show = np.zeros((inimg.shape[0], inimg.shape[1], inimg.shape[2]*2))
    #img_show[:,:,:inimg.shape[2]] = inimg
    #img_show[:,:outimg.shape[1],inimg.shape[2]:inimg.shape[2]+outimg.shape[2]] = outimg
    outimg = out_ary[:,0:h,0:w] # trim paddings
    img_show = np.concatenate((inimg, outimg), axis=2)
    bgrpic = to_bgr(img_show).copy()
    cv2.putText(bgrpic,"input",(3,15),cv2.FONT_HERSHEY_DUPLEX, 0.5,(255,0,0))
    cv2.putText(bgrpic,"output",(w+3,15),cv2.FONT_HERSHEY_DUPLEX, 0.5,(255,0,0))
    cv2.imshow("result", bgrpic)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
def main():
    parser = argparse.ArgumentParser(
        description='chainer implementation of pix2pix')
    parser.add_argument('--batchsize',
                        '-b',
                        type=int,
                        default=1,
                        help='Number of images in each mini-batch')
    parser.add_argument('--epoch',
                        '-e',
                        type=int,
                        default=200,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--gpu',
                        '-g',
                        type=int,
                        default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--dataset',
                        '-i',
                        default='./facade/base',
                        help='Directory of image files.')
    parser.add_argument('--out',
                        '-o',
                        default='result',
                        help='Directory to output the result')
    parser.add_argument('--resume',
                        '-r',
                        default='',
                        help='Resume the training from snapshot')
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    parser.add_argument('--snapshot_interval',
                        type=int,
                        default=1000,
                        help='Interval of snapshot')
    parser.add_argument('--display_interval',
                        type=int,
                        default=100,
                        help='Interval of displaying log to console')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))
    print('')

    # Set up a neural network to train
    enc = Encoder(in_ch=12)
    dec = Decoder(out_ch=3)
    dis = Discriminator(in_ch=12, out_ch=3)

    if args.gpu >= 0:
        chainer.cuda.get_device(args.gpu).use()  # Make a specified GPU current
        enc.to_gpu()  # Copy the model to the GPU
        dec.to_gpu()
        dis.to_gpu()

    # Setup an optimizer
    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec')
        return optimizer

    opt_enc = make_optimizer(enc)
    opt_dec = make_optimizer(dec)
    opt_dis = make_optimizer(dis)

    train_d = FacadeDataset(args.dataset, data_range=(1, 300))
    test_d = FacadeDataset(args.dataset, data_range=(300, 379))
    #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4)
    #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4)
    train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize)
    test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize)

    # Set up a trainer
    updater = FacadeUpdater(models=(enc, dec, dis),
                            iterator={
                                'main': train_iter,
                                'test': test_iter
                            },
                            optimizer={
                                'enc': opt_enc,
                                'dec': opt_dec,
                                'dis': opt_dis
                            },
                            device=args.gpu)
    trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)

    snapshot_interval = (args.snapshot_interval, 'iteration')
    display_interval = (args.display_interval, 'iteration')
    trainer.extend(
        extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'),
        trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        enc, 'enc_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dec, 'dec_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dis, 'dis_iter_{.updater.iteration}.npz'),
                   trigger=snapshot_interval)
    trainer.extend(extensions.LogReport(trigger=display_interval))
    trainer.extend(extensions.PrintReport([
        'epoch',
        'iteration',
        'enc/loss',
        'dec/loss',
        'dis/loss',
    ]),
                   trigger=display_interval)
    trainer.extend(extensions.ProgressBar(update_interval=10))
    trainer.extend(out_image(updater, enc, dec, 5, 5, args.seed, args.out),
                   trigger=snapshot_interval)

    if args.resume:
        # Resume from a snapshot
        chainer.serializers.load_npz(args.resume, trainer)

    # Run the training
    trainer.run()
示例#27
0
from net import Decoder
import torch
import matplotlib.pyplot as plt

input_data = torch.tensor([48.2976, -2.0373, -29.1018, 12.2312])
# input_data = torch.tensor([23.9500, -11.1840, 36.5084, -33.9963])
# input_data = torch.tensor([9.9927, -0.2395, -1.1840, -2.4094])
# input_data = torch.tensor([20.2314, -10.3747, 11.7729, -8.3415])
# input_data = torch.tensor([6.8456, 5.0325, -4.7677, -0.9231])
# input_data = torch.tensor([12.7765, -8.3813, 9.4535, -3.7266])
# input_data = torch.tensor([5.8932, 0.2757, -5.0638, 6.0578])
# input_data = torch.tensor([7.7216, 11.7913, 1.8647, -13.9440])
# input_data = torch.tensor([8.0073, -0.2398, 7.3321, -9.5089])
# input_data = torch.tensor([2.0922, 4.1933, -0.1578, -3.3220])
# input_data = (torch.randn(4))*50

print(input_data)
decoder = Decoder().eval()
model_dict = decoder.state_dict()
pretrained_dict = torch.load('model/net.pth')
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
decoder.load_state_dict(model_dict)
with torch.no_grad():
    output = decoder(input_data)
    plt.imshow(output[0], cmap='gray')
    plt.show()