コード例 #1
0
def models(modelname, cfg, word_len):
    #print(word_len)
    text_encoder = cache.get(modelname + '_text_encoder', None)
    if text_encoder is None:
        #print("text_encoder not cached")
        text_encoder = RNN_ENCODER(word_len, nhidden=cfg.TEXT.EMBEDDING_DIM)
        state_dict = torch.load(cfg.TRAIN.NET_E,
                                map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        if cfg.CUDA:
            text_encoder.cuda()
        text_encoder.eval()
        cache[modelname + '_text_encoder'] = text_encoder

    netG = cache.get(modelname + '_netG', None)
    if netG is None:
        #print("netG not cached")
        netG = G_NET()
        state_dict = torch.load(cfg.TRAIN.NET_G,
                                map_location=lambda storage, loc: storage)
        netG.load_state_dict(state_dict)
        if cfg.CUDA:
            netG.cuda()
        netG.eval()
        cache[modelname + '_netG'] = netG

    return text_encoder, netG
コード例 #2
0
ファイル: eval.py プロジェクト: jennifergreb/Marrow
def models(word_len):
    print('Loading Model', word_len)
    text_encoder = cache.get('text_encoder')
    print('Text enconder', text_encoder)
    if text_encoder is None:
        print("text_encoder not cached")
        text_encoder = RNN_ENCODER(word_len, nhidden=256)
        state_dict = torch.load('../DAMSMencoders/coco/text_encoder100.pth', map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        print('loaded text encoder')
        text_encoder.cuda()
        print('text encoder cuda')
        text_encoder.eval()
        print('text encoder eval')
        #cache.set('text_encoder', text_encoder, timeout=60 * 60 * 24)

    print('Got Text Encoder, moving to netG')
    netG = cache.get('netG')
    if netG is None:
        print("netG not cached")
        netG = G_NET()
        state_dict = torch.load('../models/coco_AttnGAN2.pth', map_location=lambda storage, loc: storage)
        netG.load_state_dict(state_dict)
        if cfg.CUDA:
            netG.cuda()
        netG.eval()
        #cache.set('netG', netG, timeout=60 * 60 * 24)
    print('Got NetG')
    return text_encoder, netG
コード例 #3
0
ファイル: trainer.py プロジェクト: iihihiuh/origami_inference
    def load_network_stageII(self):
        from model import G_NET, D_NET

        netG = G_NET()
        netG.apply(weights_init)
        print(netG)
        if cfg.NET_G != '':
            state_dict = \
                torch.load(cfg.NET_G,
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load from: ', cfg.NET_G)

        netD = D_NET()
        netD.apply(weights_init)
        if cfg.NET_D != '':
            state_dict = \
                torch.load(cfg.NET_D,
                           map_location=lambda storage, loc: storage)
            netD.load_state_dict(state_dict)
            print('Load from: ', cfg.NET_D)
        print(netD)

        if cfg.CUDA:
            netG.cuda()
            netD.cuda()
        return netG, netD
コード例 #4
0
ファイル: test.py プロジェクト: ryanbrand/GANTextToImage
def test(cfg_file, embedding_t7_path):

    cfg_from_file(cfg_file)

    cfg.GPU_ID = 0

    print('Using config:')
    pprint.pprint(cfg)

    netG = G_NET()
    netG.apply(weights_init)
    netG = torch.nn.DataParallel(netG, device_ids=[0])
    print(netG)
    #state_dict = torch.load('../gan/models/birds_3stages/netG_26000.pth')
    print(cfg.TRAIN.NET_G)
    state_path = '/home/ubuntu/GANTextToImage/gan/' + cfg.TRAIN.NET_G
    state_dict = torch.load(state_path, map_location=lambda storage, loc: storage)
    netG.load_state_dict(state_dict)
    # print('Load ', '../gan/models/flowers_1stage/netG_6000.pth')

    # the path to save generated images
    s_tmp = cfg.TRAIN.NET_G
    istart = s_tmp.rfind('_') + 1
    iend = s_tmp.rfind('.')
    iteration = int(s_tmp[istart:iend])
    s_tmp = s_tmp[:s_tmp.rfind('/')]
    save_dir = '%s/iteration%d' % (s_tmp, iteration)

    nz = cfg.GAN.Z_DIM
    noise = Variable(torch.FloatTensor(1, nz))
    if cfg.CUDA:
        netG.cuda()
        noise = noise.cuda()

    # switch to evaluate mode
    netG.eval()

    t_embedding = load_lua(embedding_t7_path)
    t_embedding = t_embedding.unsqueeze(0)
    print(t_embedding.size())   

    if cfg.CUDA:
        t_embedding = Variable(t_embedding).cuda()
    else:
        t_embedding = Variable(t_embedding)
    # print(t_embeddings[:, 0, :], t_embeddings.size(1))

    embedding_dim = t_embedding.size(1)
    noise.data.resize_(1, nz)
    noise.data.normal_(0, 1)

    fake_img_list = []
    # for i in range(embedding_dim):
    fake_imgs, _, _ = netG(noise, t_embedding[:, 0, :])
    if cfg.TEST.B_EXAMPLE:
            # fake_img_list.append(fake_imgs[0].data.cpu())
            # fake_img_list.append(fake_imgs[1].data.cpu())
        fake_img_list.append(fake_imgs[2].data.cpu())
    else:
        save_singleimages(fake_imgs[-1], '/home/ubuntu/GANTextToImage/static', 0, 256)
コード例 #5
0
ファイル: eval.py プロジェクト: OnizukaLab/AttnGAN
def models(word_len):
    #print(word_len)
    text_encoder = cache.get('text_encoder')
    if text_encoder is None:
        #print("text_encoder not cached")
        text_encoder = RNN_ENCODER(word_len, nhidden=cfg.TEXT.EMBEDDING_DIM)
        state_dict = torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        if cfg.CUDA:
            text_encoder.cuda()
        text_encoder.eval()
        cache.set('text_encoder', text_encoder, timeout=60 * 60 * 24)

    netG = cache.get('netG')
    if netG is None:
        #print("netG not cached")
        netG = G_NET()
        state_dict = torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
        netG.load_state_dict(state_dict)
        if cfg.CUDA:
            netG.cuda()
        netG.eval()
        cache.set('netG', netG, timeout=60 * 60 * 24)

    return text_encoder, netG
コード例 #6
0
ファイル: trainer.py プロジェクト: synsypa/Recipe2ImageGAN
def load_network(gpus):
    netG = G_NET()
    netG.apply(weights_init)
    netG = torch.nn.DataParallel(netG, device_ids=gpus)
    print(netG)

    netsD = []
    if cfg.TREE.BRANCH_NUM > 0:
        netsD.append(D_NET64())
    if cfg.TREE.BRANCH_NUM > 1:
        netsD.append(D_NET128())
    if cfg.TREE.BRANCH_NUM > 2:
        netsD.append(D_NET256())
    if cfg.TREE.BRANCH_NUM > 3:
        netsD.append(D_NET512())
    if cfg.TREE.BRANCH_NUM > 4:
        netsD.append(D_NET1024())
    # TODO: if cfg.TREE.BRANCH_NUM > 5:

    for i in range(len(netsD)):
        netsD[i].apply(weights_init)
        netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
        # print(netsD[i])
    print('# of netsD', len(netsD))

    count = 0
    if cfg.TRAIN.NET_G != '':
        state_dict = torch.load(cfg.TRAIN.NET_G)
        netG.load_state_dict(state_dict)
        print('Load ', cfg.TRAIN.NET_G)

        try:
            istart = cfg.TRAIN.NET_G.rfind('_') + 1
            iend = cfg.TRAIN.NET_G.rfind('.')
            count = cfg.TRAIN.NET_G[istart:iend]
            count = int(count)
        except:
            last_run_dir = cfg.DATA_DIR + '/' + cfg.LAST_RUN_DIR + '/Model'
            with open(last_run_dir + '/count.txt', 'r') as f:
                count = int(f.read())

        count = int(count) + 1

    if cfg.TRAIN.NET_D != '':
        for i in range(len(netsD)):
            print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
            state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
            netsD[i].load_state_dict(state_dict)

    inception_model = INCEPTION_V3()

    if cfg.CUDA:
        netG.cuda()
        for i in range(len(netsD)):
            netsD[i].cuda()
        inception_model = inception_model.cuda()
    inception_model.eval()

    return netG, netsD, len(netsD), inception_model, count
コード例 #7
0
def load_network(gpus, dictionary=None):
    netG = G_NET()
    netG.apply(weights_init)
    netG = torch.nn.DataParallel(netG, device_ids=gpus)
    print(netG)

    netsD = []
    if cfg.TREE.BRANCH_NUM > 0:
        netsD.append(D_NET64())
    if cfg.TREE.BRANCH_NUM > 1:
        netsD.append(D_NET128())
    if cfg.TREE.BRANCH_NUM > 2:
        netsD.append(D_NET256())
    if cfg.TREE.BRANCH_NUM > 3:
        netsD.append(D_NET512())
    if cfg.TREE.BRANCH_NUM > 4:
        netsD.append(D_NET1024())
    # TODO: if cfg.TREE.BRANCH_NUM > 5:

    for i in range(len(netsD)):
        netsD[i].apply(weights_init)
        netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
        # print(netsD[i])
    print('# of netsD', len(netsD))

    count = 0
    if cfg.TRAIN.NET_G != '':
        state_dict = torch.load(cfg.TRAIN.NET_G)
        netG.load_state_dict(state_dict)
        print('Load ', cfg.TRAIN.NET_G)

        istart = cfg.TRAIN.NET_G.rfind('_') + 1
        iend = cfg.TRAIN.NET_G.rfind('.')
        count = cfg.TRAIN.NET_G[istart:iend]
        count = int(count) + 1

    if cfg.TRAIN.NET_D != '':
        for i in range(len(netsD)):
            print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
            state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
            netsD[i].load_state_dict(state_dict)

    inception_model = INCEPTION_V3()

    embedding_model = load_embedding_model(dictionary)
    for param in embedding_model.parameters():
        param.requires_grad = False

    if cfg.CUDA:
        netG.cuda()
        for i in range(len(netsD)):
            netsD[i].cuda()
        inception_model = inception_model.cuda()
        embedding_model.cuda()

    inception_model.eval()

    return netG, netsD, len(netsD), inception_model, embedding_model, count
コード例 #8
0
ファイル: trainer.py プロジェクト: AustinXY/fineGAN
    def evaluate_finegan(self):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for model not found!')
        else:
            # Build and load the generator
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            model_dict = netG.state_dict()

            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)['birds128']

            state_dict = {k: v for k, v in state_dict.items() if k in model_dict}

            model_dict.update(state_dict)
            netG.load_state_dict(model_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # Uncomment this to print Generator layers
            print(netG)

            sys

            nz = cfg.GAN.Z_DIM
            noise = torch.FloatTensor(self.batch_size, nz)
            noise.data.normal_(0, 1)

            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            netG.eval()

            background_class = cfg.TEST_BACKGROUND_CLASS
            parent_class = cfg.TEST_PARENT_CLASS
            child_class = cfg.TEST_CHILD_CLASS
            bg_code = torch.zeros([self.batch_size, cfg.FINE_GRAINED_CATEGORIES])
            p_code = torch.zeros([self.batch_size, cfg.SUPER_CATEGORIES])
            c_code = torch.zeros([self.batch_size, cfg.FINE_GRAINED_CATEGORIES])

            for j in range(self.batch_size):
                bg_code[j][background_class] = 1
                p_code[j][parent_class] = 1
                c_code[j][child_class] = 1

            fake_imgs, fg_imgs, mk_imgs, fgmk_imgs = netG(noise, c_code, p_code, bg_code) # Forward pass through the generator

            self.save_image(fake_imgs[0][0], self.save_dir, 'background')
            self.save_image(fake_imgs[1][0], self.save_dir, 'parent_final')
            self.save_image(fake_imgs[2][0], self.save_dir, 'child_final')
            self.save_image(fg_imgs[0][0], self.save_dir, 'parent_foreground')
            self.save_image(fg_imgs[1][0], self.save_dir, 'child_foreground')
            self.save_image(mk_imgs[0][0], self.save_dir, 'parent_mask')
            self.save_image(mk_imgs[1][0], self.save_dir, 'child_mask')
            self.save_image(fgmk_imgs[0][0], self.save_dir, 'parent_foreground_masked')
            self.save_image(fgmk_imgs[1][0], self.save_dir, 'child_foreground_masked')
コード例 #9
0
ファイル: trainer.py プロジェクト: BoChenGroup/VHE-GAN
def load_network(gpus):
    netEn_img = MLP_ENCODER_IMG()
    netEn_img.apply(weights_init)
    netEn_img = torch.nn.DataParallel(netEn_img, device_ids=gpus)
    print(netEn_img)

    netG = G_NET()
    netG.apply(weights_init)
    netG = torch.nn.DataParallel(netG, device_ids=gpus)
    print(netG)

    netsD = []
    if cfg.TREE.BRANCH_NUM > 0:
        netsD.append(D_NET64())
    if cfg.TREE.BRANCH_NUM > 1:
        netsD.append(D_NET128())
    if cfg.TREE.BRANCH_NUM > 2:
        netsD.append(D_NET256())

    for i in xrange(len(netsD)):
        netsD[i].apply(weights_init)
        netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
    print('# of netsD', len(netsD))

    count = 0
    if cfg.TRAIN.NET_G != '':
        state_dict = torch.load(cfg.TRAIN.NET_G)
        netG.load_state_dict(state_dict)
        print('Load ', cfg.TRAIN.NET_G)

        istart = cfg.TRAIN.NET_G.rfind('_') + 1
        iend = cfg.TRAIN.NET_G.rfind('.')
        count = cfg.TRAIN.NET_G[istart:iend]
        count = int(count) + 1

    if cfg.TRAIN.NET_D != '':
        for i in xrange(len(netsD)):
            print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
            state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
            netsD[i].load_state_dict(state_dict)

    if cfg.TRAIN.NET_MLP_IMG != '':
        state_dict = torch.load(cfg.TRAIN.NET_MLP_IMG)
        netEn_img.load_state_dict(state_dict)
        print('Load ', cfg.TRAIN.NET_MLP_IMG)

    inception_model = INCEPTION_V3()

    if cfg.CUDA:
        netG.cuda()
        netEn_img = netEn_img.cuda()
        for i in xrange(len(netsD)):
            netsD[i].cuda()
        inception_model = inception_model.cuda()
    inception_model.eval()

    return netG, netsD, netEn_img, inception_model, len(netsD), count
コード例 #10
0
ファイル: trainer.py プロジェクト: loveisessential/showcase
    def evaluate(self, split_dir):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            # Build and load the generator
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            print(netG)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # the path to save generated images
            s_tmp = cfg.TRAIN.NET_G
            istart = s_tmp.rfind('_') + 1
            iend = s_tmp.rfind('.')
            iteration = int(s_tmp[istart:iend])
            s_tmp = s_tmp[:s_tmp.rfind('/')]
            save_dir = '%s/iteration%d/%s' % (s_tmp, iteration, split_dir)
            if cfg.TEST.B_EXAMPLE:
                folder = '%s/super' % (save_dir)
            else:
                folder = '%s/single' % (save_dir)
            print('Make a new folder: ', folder)
            mkdir_p(folder)

            nz = cfg.GAN.Z_DIM
            noise = Variable(torch.FloatTensor(self.batch_size, nz))
            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            # switch to evaluate mode
            netG.eval()
            num_batches = int(cfg.TEST.SAMPLE_NUM / self.batch_size)
            cnt = 0
            for step in range(num_batches):
                noise.data.normal_(0, 1)
                #hmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxstart
                fake_imgs, layers_output, _, _ = netG(noise)
                if len(layers_output) != len(lamdas):
                    print("please check lamdas length")
                #hmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxend
                if cfg.TEST.B_EXAMPLE:
                    self.save_superimages(fake_imgs[-1], folder, cnt, 256)
                else:
                    self.save_singleimages(fake_imgs[-1], folder, cnt, 256)
                    # self.save_singleimages(fake_imgs[-2], folder, 128)
                    # self.save_singleimages(fake_imgs[-3], folder, 64)
                cnt += self.batch_size
コード例 #11
0
ファイル: trainer.py プロジェクト: tensoralex/StackGAN-v2
def load_network(gpus):
    netG = G_NET()
    netG.apply(weights_init)
    netG = torch.nn.DataParallel(netG, device_ids=gpus)
    print(netG)

    netsD = []
    if cfg.TREE.BRANCH_NUM > 0:
        netsD.append(D_NET64())
    if cfg.TREE.BRANCH_NUM > 1:
        netsD.append(D_NET128())
    if cfg.TREE.BRANCH_NUM > 2:
        netsD.append(D_NET256())
    if cfg.TREE.BRANCH_NUM > 3:
        netsD.append(D_NET512())
    if cfg.TREE.BRANCH_NUM > 4:
        netsD.append(D_NET1024())
    # TODO: if cfg.TREE.BRANCH_NUM > 5:

    for i in range(len(netsD)):
        netsD[i].apply(weights_init)
        netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
        # print(netsD[i])
    print('# of netsD', len(netsD))

    count = 0
    if cfg.TRAIN.NET_G != '':
        state_dict = torch.load(cfg.TRAIN.NET_G)
        netG.load_state_dict(state_dict)
        print('Load ', cfg.TRAIN.NET_G)

        istart = cfg.TRAIN.NET_G.rfind('_') + 1
        iend = cfg.TRAIN.NET_G.rfind('.')
        count = cfg.TRAIN.NET_G[istart:iend]
        count = int(count) + 1

    if cfg.TRAIN.NET_D != '':
        for i in range(len(netsD)):
            print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
            state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
            netsD[i].load_state_dict(state_dict)

    inception_model = INCEPTION_V3()

    if cfg.CUDA:
        netG.cuda()
        for i in range(len(netsD)):
            netsD[i].cuda()
        inception_model = inception_model.cuda()
    inception_model.eval()

    return netG, netsD, len(netsD), inception_model, count
コード例 #12
0
def load_network(gpus):
    netG = G_NET()
    netG.apply(weights_init)
    netG = torch.nn.DataParallel(netG, device_ids=gpus)
    print(netG)

    netsD = []
    # 128 * 128
    if cfg.TREE.BRANCH_NUM > 1:
        for i in range(
                3):  # 3 discriminators for background, parent and child stage
            netsD.append(D_NET128(i))

    # 256 * 256
    if cfg.TREE.BRANCH_NUM > 2:
        for i in range(
                3):  # 3 discriminators for background, parent and child stage
            netsD.append(D_NET256(i))

    # for i in range(3): # 3 discriminators for background, parent and child stage
    #     netsD.append(D_NET128(i))

    for i in range(len(netsD)):
        netsD[i].apply(weights_init)
        netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
        print(netsD[i])

    count = 0

    if cfg.TRAIN.NET_G != '':
        state_dict = torch.load(cfg.TRAIN.NET_G)
        netG.load_state_dict(state_dict)
        print('Load ', cfg.TRAIN.NET_G)

        istart = cfg.TRAIN.NET_G.rfind('_') + 1
        iend = cfg.TRAIN.NET_G.rfind('.')
        count = cfg.TRAIN.NET_G[istart:iend]
        count = int(count) + 1

    if cfg.TRAIN.NET_D != '':
        for i in range(len(netsD)):
            print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
            state_dict = torch.load('%s_%d.pth' % (cfg.TRAIN.NET_D, i))
            netsD[i].load_state_dict(state_dict)

    if cfg.CUDA:
        netG.cuda()
        for i in range(len(netsD)):
            netsD[i].cuda()

    return netG, netsD, len(netsD), count
コード例 #13
0
ファイル: trainer.py プロジェクト: zxs789/Obj-GAN
    def build_models(self):
        netG = G_NET(len(self.cats_index_dict))
        netINSD = INS_D_NET(len(self.cats_index_dict))
        netGLBD = GLB_D_NET(len(self.cats_index_dict))

        netG.apply(weights_init)
        netINSD.apply(weights_init)
        netGLBD.apply(weights_init)

        if cfg.CUDA:
            netG.cuda()
            netINSD.cuda()
            netGLBD.cuda()

            if len(cfg.GPU_IDS) > 1:
                netG = nn.DataParallel(netG)
                netG.to(self.device)
                netINSD = nn.DataParallel(netINSD)
                netINSD.to(self.device)
                netGLBD = nn.DataParallel(netGLBD)
                netGLBD.to(self.device)

        # ########################################################### #
        epoch = 0
        if cfg.TRAIN.NET_G != '':
            state_dict = \
                torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load G from: ', cfg.TRAIN.NET_G)
            filename = path_leaf(cfg.TRAIN.NET_G)
            istart = filename.rfind('_') + 1
            iend = filename.rfind('.')
            epoch = filename[istart:iend]
            epoch = int(epoch) + 1

            Gname = cfg.TRAIN.NET_G
            s_tmp = Gname[:Gname.rfind('/')]
            Dname = '%s/netINSD.pth' % (s_tmp)
            print('Load INSD from: ', Dname)
            state_dict = \
                torch.load(Dname, map_location=lambda storage, loc: storage)
            netINSD.load_state_dict(state_dict)

            s_tmp = Gname[:Gname.rfind('/')]
            Dname = '%s/netGLBD.pth' % (s_tmp)
            print('Load GLBD from: ', Dname)
            state_dict = \
                torch.load(Dname, map_location=lambda storage, loc: storage)
            netGLBD.load_state_dict(state_dict)

        return [netG, netINSD, netGLBD, epoch]
コード例 #14
0
ファイル: trainer.py プロジェクト: tensoralex/StackGAN-v2
    def evaluate(self, split_dir):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            # Build and load the generator
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            print(netG)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # the path to save generated images
            s_tmp = cfg.TRAIN.NET_G
            istart = s_tmp.rfind('_') + 1
            iend = s_tmp.rfind('.')
            iteration = int(s_tmp[istart:iend])
            s_tmp = s_tmp[:s_tmp.rfind('/')]
            save_dir = '%s/iteration%d/%s' % (s_tmp, iteration, split_dir)
            if cfg.TEST.B_EXAMPLE:
                folder = '%s/super' % (save_dir)
            else:
                folder = '%s/single' % (save_dir)
            print('Make a new folder: ', folder)
            mkdir_p(folder)

            nz = cfg.GAN.Z_DIM
            noise = Variable(torch.FloatTensor(self.batch_size, nz))
            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            # switch to evaluate mode
            netG.eval()
            num_batches = int(cfg.TEST.SAMPLE_NUM / self.batch_size)
            cnt = 0
            for step in xrange(num_batches):
                noise.data.normal_(0, 1)
                fake_imgs, _, _ = netG(noise)
                if cfg.TEST.B_EXAMPLE:
                    self.save_superimages(fake_imgs[-1], folder, cnt, 256)
                else:
                    self.save_singleimages(fake_imgs[-1], folder, cnt, 256)
                    # self.save_singleimages(fake_imgs[-2], folder, 128)
                    # self.save_singleimages(fake_imgs[-3], folder, 64)
                cnt += self.batch_size
コード例 #15
0
def load_network(gpus):
    netG = G_NET(start_depth)
    netG.apply(weights_init)
    netG = torch.nn.DataParallel(netG, device_ids=gpus)
    print(netG)

    netsD = []
    netsD.append(D_NET_BG(start_depth))
    netsD.append(D_NET_PC(1, start_depth))
    netsD.append(D_NET_PC(2, start_depth))
    netsD.append(D_NET_BG_PG(start_depth))

    for i in range(len(netsD)):
        netsD[i].apply(weights_init)
        netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
        print(netsD[i])

    count = 0

    if cfg.TRAIN.NET_G != '':
        state_dict = torch.load(cfg.TRAIN.NET_G)
        netG.load_state_dict(state_dict)
        print('Load ', cfg.TRAIN.NET_G)

        istart = cfg.TRAIN.NET_G.rfind('netG_') + 5
        iend = cfg.TRAIN.NET_G.rfind('_depth')
        count = cfg.TRAIN.NET_G[istart:iend]
        count = int(count)
        istart = cfg.TRAIN.NET_G.rfind('depth')
        iend = cfg.TRAIN.NET_G.rfind('.')
        _depth = cfg.TRAIN.NET_G[istart:iend]

    if cfg.TRAIN.NET_D != '':
        for i in range(len(netsD)):
            print('Load %s%d_%s.pth' % (cfg.TRAIN.NET_D, i, _depth))
            state_dict = torch.load('%s%d_%s.pth' %
                                    (cfg.TRAIN.NET_D, i, _depth))
            netsD[i].load_state_dict(state_dict)

    if cfg.CUDA:
        netG.cuda()
        for i in range(len(netsD)):
            netsD[i].cuda()

    return netG, netsD, len(netsD), count
コード例 #16
0
def models(word_len):
    text_encoder = cache.get('text_encoder')
    if text_encoder is None:
        text_encoder = RNN_ENCODER(word_len, nhidden=256)
        state_dict = torch.load('../DAMSMencoders/coco/text_encoder100.pth', map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        text_encoder.cuda()
        text_encoder.eval()
        #cache.set('text_encoder', text_encoder, timeout=60 * 60 * 24)

    netG = cache.get('netG')
    if netG is None:
        netG = G_NET()
        state_dict = torch.load('../models/coco_AttnGAN2.pth', map_location=lambda storage, loc: storage)
        netG.load_state_dict(state_dict)
        if cfg.CUDA:
            netG.cuda()
        netG.eval()
        #cache.set('netG', netG, timeout=60 * 60 * 24)
    return text_encoder, netG
コード例 #17
0
    def evaluate(self, split_dir):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            # Build and load the generator
            if split_dir == 'test':
                split_dir = 'valid'
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            print(netG)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load('/content/drive/My Drive/Colab Notebooks/StackGAN-v2-master/models/netG_210000.pth',
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # the path to save generated images
            s_tmp = cfg.TRAIN.NET_G
            istart = s_tmp.rfind('_') + 1
            iend = s_tmp.rfind('.')
            iteration = int(s_tmp[istart:iend])
            s_tmp = s_tmp[:s_tmp.rfind('/')]
            save_dir = '%s/iteration%d' % (s_tmp, iteration)

            nz = cfg.GAN.Z_DIM
            noise = Variable(torch.FloatTensor(self.batch_size, nz))
            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            # switch to evaluate mode
            netG.eval()
            for step, data in enumerate(self.data_loader, 0):
                imgs, t_embeddings, filenames = data
                if cfg.CUDA:
                    t_embeddings = Variable(t_embeddings).cuda()
                else:
                    t_embeddings = Variable(t_embeddings)
                # print(t_embeddings[:, 0, :], t_embeddings.size(1))

                embedding_dim = t_embeddings.size(1)
                batch_size = imgs[0].size(0)
                noise.data.resize_(batch_size, nz)
                noise.data.normal_(0, 1)

                fake_img_list = []
                for i in range(embedding_dim):
                    fake_imgs, _, _ = netG(noise, t_embeddings[:, i, :])
                    if cfg.TEST.B_EXAMPLE:
                        # fake_img_list.append(fake_imgs[0].data.cpu())
                        # fake_img_list.append(fake_imgs[1].data.cpu())
                        fake_img_list.append(fake_imgs[2].data.cpu())
                    else:
                        self.save_singleimages(fake_imgs[-1], filenames,
                                               save_dir, split_dir, i, 256)
                        # self.save_singleimages(fake_imgs[-2], filenames,
                        #                        save_dir, split_dir, i, 128)
                        # self.save_singleimages(fake_imgs[-3], filenames,
                        #                        save_dir, split_dir, i, 64)
                    # break
                if cfg.TEST.B_EXAMPLE:
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 64)
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 128)
                    self.save_superimages(fake_img_list, filenames,
                                          save_dir, split_dir, 256)
コード例 #18
0
    def evaluate(self, split_dir):
        inception_model = INCEPTION_V3()
        # fid_model = FID_INCEPTION()
        if cfg.CUDA:
            inception_model.cuda()
        #     fid_model.cuda()
        inception_model.eval()
        # fid_model.eval()

        if cfg.TRAIN.NET_G == '':
            print('Error: the path for models is not found!')
        else:
            # Build and load the generator
            if split_dir == 'test':
                split_dir = 'valid'
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            # print(netG)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # the path to save generated images
            # s_tmp = cfg.TRAIN.NET_G
            # istart = s_tmp.rfind('_') + 1
            # iend = s_tmp.rfind('.')
            # iteration = int(s_tmp[istart:iend])
            # s_tmp = s_tmp[:s_tmp.rfind('/')]
            # save_dir = '%s/iteration%d' % (s_tmp, iteration)
            # save_dir = 'C:\\Users\\alper\\PycharmProjects\\MSGAN\\StackGAN++-Mode-Seeking\\results'
            save_dir = "D:\\results"

            nz = cfg.GAN.Z_DIM
            n_samples = 50
            # noise = Variable(torch.FloatTensor(self.batch_size, nz))
            noise = Variable(torch.FloatTensor(n_samples, nz))
            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            # switch to evaluate mode
            netG.eval()
            for step, data in enumerate(tqdm(self.data_loader)):
                # if step == 8:
                #     break
                imgs, t_embeddings, filenames = data
                if cfg.CUDA:
                    t_embeddings = Variable(t_embeddings).cuda()
                else:
                    t_embeddings = Variable(t_embeddings)
                # print(t_embeddings[:, 0, :], t_embeddings.size(1))

                embedding_dim = t_embeddings.size(1)
                # batch_size = imgs[0].size(0)
                # noise.data.resize_(batch_size, nz)
                noise.data.normal_(0, 1)

                fake_img_list = []
                inception_score_list = []
                fid_list = []
                score_list = []
                predictions = []
                fids = []
                for i in range(embedding_dim):
                    inception_score_list.append([])
                    fid_list.append([])
                    score_list.append([])

                    emb_imgs = []
                    for j in range(n_samples):
                        noise_j = noise[j].unsqueeze(0)
                        t_embeddings_i = t_embeddings[:, i, :]
                        fake_imgs, _, _ = netG(noise_j, t_embeddings_i)
                        # filenames_number ='_sample_%2.2d'%(j)
                        # filenames_new = []
                        # filenames_new.append(filenames[-1]+filenames_number)
                        # filenames_new = tuple(filenames_new)

                        # for selecting reasonable images
                        pred = inception_model(fake_imgs[-1].detach())
                        pred = pred.data.cpu().numpy()
                        predictions.append(pred)
                        bird_indices = [
                            7, 8, 9, 10, 11, 13, 15, 16, 17, 18, 19, 21, 23,
                            81, 84, 85, 86, 88, 90, 91, 93, 94, 95, 96, 97, 99,
                            129, 130, 133, 134, 135, 138, 141, 142, 143, 144,
                            146, 517
                        ]
                        score = np.max(pred[0, bird_indices])
                        score_list[i].append((j, score))
                        emb_imgs.append(fake_imgs[2].data.cpu())
                        if cfg.TEST.B_EXAMPLE:
                            # fake_img_list.append(fake_imgs[0].data.cpu())
                            # fake_img_list.append(fake_imgs[1].data.cpu())
                            fake_img_list.append(fake_imgs[2].data.cpu())
                        else:
                            self.save_singleimages(fake_imgs[-1], filenames, j,
                                                   save_dir, split_dir, i, 256)
                        # self.save_singleimages(fake_imgs[-2], filenames,
                        #                        save_dir, split_dir, i, 128)
                        # self.save_singleimages(fake_imgs[-3], filenames,
                        #                        save_dir, split_dir, i, 64)
                    # break
                    score_list[i] = sorted(score_list[i],
                                           key=lambda x: x[1],
                                           reverse=True)[:5]
                    # for FID score
                    # ffi = [i[0].numpy() for i in emb_imgs]
                    fake_filtered_images = [
                        fake_img_list[i][0].numpy()
                        for i in range(len(fake_img_list))
                    ]
                    img_dir = os.path.join(cfg.DATA_DIR, "CUB_200_2011",
                                           "images",
                                           filenames[0].split("/")[0])
                    img_files = [
                        os.path.join(img_dir, i) for i in os.listdir(img_dir)
                    ]

                    # act_real = get_activations(img_files, fid_model)
                    # mu_real, sigma_real = get_fid_stats(act_real)
                    # print("mu_real: {}, sigma_real: {}".format(mu_real, sigma_real))

                    np_imgs = np.array(fake_filtered_images)
                    # print(np_imgs.shape)

                    # # print(type(np_imgs[0]))
                    # act_fake = get_activations(np_imgs, fid_model, img=True)
                    # mu_fake, sigma_fake = get_fid_stats(act_fake)
                    # fid_score = frechet_distance(mu_real, sigma_real, mu_fake, sigma_fake)
                    # fids.append(fid_score)
                    # print("mu_fake: {}, sigma_fake: {}".format(mu_fake, sigma_fake))
                # print(inception_score_list)

                # # calculate inception score
                # predictions = np.concatenate(predictions, 0)
                # mean, std = compute_inception_score(predictions, 10)
                # mean_nlpp, std_nlpp = \
                #     negative_log_posterior_probability(predictions, 10)
                # inception_score_list.append((mean, std, mean_nlpp, std_nlpp))

                # # for FID score
                # fake_filtered_images = [fake_img_list[i*n_samples + k[0]][0].numpy() for i, j in enumerate(score_list) for k in j]
                # # fake_filtered_images = [fake_img_list[i][0].numpy() for i in range(len(fake_img_list))]
                # img_dir = os.path.join(cfg.DATA_DIR, "CUB_200_2011", "images", filenames[0].split("/")[0])
                # img_files = [os.path.join(img_dir, i) for i in os.listdir(img_dir)]
                #
                # act_real = get_activations(img_files, fid_model)
                # mu_real, sigma_real = get_fid_stats(act_real)
                # # print("mu_real: {}, sigma_real: {}".format(mu_real, sigma_real))
                #
                # np_imgs = np.array(fake_filtered_images)
                # # print(np_imgs.shape)
                #
                # # print(type(np_imgs[0]))
                # act_fake = get_activations(np_imgs, fid_model, img=True)
                # mu_fake, sigma_fake = get_fid_stats(act_fake)
                # # print("mu_fake: {}, sigma_fake: {}".format(mu_fake, sigma_fake))
                #
                # # fid_score = frechet_distance(mu_real, sigma_real, mu_fake, sigma_fake)
                # fid_score = np.mean(fids)
                # fid_list.append(fid_score)
                # stats = 'step: {}, FID: {}, inception_score: {}, nlpp: {}\n'.format(step, fid_score, (mean, std), (mean_nlpp, std_nlpp))
                # with open("results\\stats.txt", "a+") as f:
                #     f.write(stats)
                # print(stats)

                if cfg.TEST.B_EXAMPLE:
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 64)
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 128)
                    if cfg.TEST.FILTER:
                        images_to_save = [
                            fake_img_list[i * n_samples + k[0]]
                            for i, j in enumerate(score_list) for k in j
                        ]
                    else:
                        images_to_save = fake_img_list
                    self.save_superimages(images_to_save, filenames, save_dir,
                                          split_dir, 256)
コード例 #19
0
    def gen_samples(self, idx):

        text_encoder = RNN_ENCODER(self.n_words,
                                   nhidden=cfg.TEXT.EMBEDDING_DIM)
        state_dict = torch.load(cfg.TRAIN.NET_E,
                                map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        print('Load text encoder from: {}'.format(cfg.TRAIN.NET_E))
        text_encoder = text_encoder.cuda()
        text_encoder.eval()

        netG = G_NET()
        state_dict = torch.load(cfg.TRAIN.NET_G,
                                map_location=lambda storage, loc: storage)
        netG.load_state_dict(state_dict)
        print('Load G from: {}'.format(cfg.TRAIN.NET_G))
        netG.cuda()
        netG.eval()

        s_tmp = cfg.TRAIN.NET_G[:cfg.TRAIN.NET_G.rfind('.pth')]
        save_dir = '%s/samples' % (s_tmp)
        mkdir_p(save_dir)

        batch_size = self.batch_size
        nz = cfg.GAN.Z_DIM
        with torch.no_grad():
            noise = Variable(torch.FloatTensor(batch_size, nz))
        noise = noise.cuda()

        step = 0
        data_iter = iter(self.data_loader)

        while step < self.num_batches:
            data = data_iter.next()
            imgs, captions, cap_lens, class_ids, sorted_cap_indices = self.prepare_data(
                data)
            hidden = text_encoder.init_hidden(batch_size)
            words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
            mask = (captions == 0)
            num_words = words_embs.size(2)
            if mask.size(1) > num_words:
                mask = mask[:, :num_words]
            for i in range(10):
                noise.data.normal_(0, 1)
                fake_imgs, attention_maps, _, _ = netG(noise, sent_emb,
                                                       words_embs, mask)
                cap_lens_np = cap_lens.cpu().data.numpy()
                for j in range(batch_size):
                    right_idx = step * batch_size + sorted_cap_indices[j]
                    save_name = '%s/%d_s_%d' % (save_dir, i, right_idx)
                    original_idx = idx[right_idx]
                    shutil.copyfile(
                        '/.local/AttnGAN/data/FashionSynthesis/test/original/test128_{}.png'
                        .format(original_idx + 1),
                        save_dir + '/test128_{0}_{1}.png'.format(
                            original_idx + 1, right_idx))
                    for k in range(len(fake_imgs)):
                        im = fake_imgs[k][j].data.cpu().numpy()
                        im = (im + 1.0) * 127.5
                        im = im.astype(np.uint8)
                        im = np.transpose(im, (1, 2, 0))
                        im = Image.fromarray(im)
                        fullpath = '%s_g%d.png' % (save_name, k)
                        im.save(fullpath)
                    for k in range(len(attention_maps)):
                        if len(fake_imgs) > 1:
                            im = fake_imgs[k + 1].detach().cpu()
                        else:
                            im = fake_imgs[0].detach().cpu()
                        attn_maps = attention_maps[k]
                        att_sze = attn_maps.size(2)
                        img_set, sentences = \
                            build_super_images2(im[j].unsqueeze(0),
                                                captions[j].unsqueeze(0),
                                                [cap_lens_np[j]], self.ixtoword,
                                                [attn_maps[j]], att_sze)
                        if img_set is not None:
                            im = Image.fromarray(img_set)
                            fullpath = '%s_a%d.png' % (save_name, k)
                            im.save(fullpath)
            step += 1
コード例 #20
0
ファイル: trainer.py プロジェクト: loveisessential/showcase
    def evaluate(self, split_dir):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            # Build and load the generator
            if split_dir == 'test':
                split_dir = 'valid'
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            print(netG)
            #state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)

            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # the path to save generated images
            # s_tmp = cfg.TRAIN.NET_G
            # istart = s_tmp.rfind('_') + 1
            # iend = s_tmp.rfind('.')
            # iteration = int(s_tmp[istart:iend])
            # s_tmp = s_tmp[:s_tmp.rfind('/')]
            # save_dir = '%s/iteration%d' % (s_tmp, iteration)
            save_dir = 'J:\\qimao\\Text-to-image\\results\\Plugin-v1-210K-random50'

            nz = cfg.GAN.Z_DIM
            n_samples = 50
            # noise = Variable(torch.FloatTensor(self.batch_size, nz))
            noise = Variable(torch.FloatTensor(n_samples, nz))
            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            # switch to evaluate mode
            netG.eval()
            for step, data in enumerate(self.data_loader, 0):
                imgs, t_embeddings, filenames = data
                if cfg.CUDA:
                    t_embeddings = Variable(t_embeddings).cuda()
                else:
                    t_embeddings = Variable(t_embeddings)
                # print(t_embeddings[:, 0, :], t_embeddings.size(1))

                embedding_dim = t_embeddings.size(1)
                # batch_size = imgs[0].size(0)
                # noise.data.resize_(batch_size, nz)
                noise.data.normal_(0, 1)

                fake_img_list = []
                for i in range(embedding_dim):
                    for j in range(n_samples):
                        noise_j = noise[j].unsqueeze(0)
                        t_embeddings_i = t_embeddings[:, i, :]

                        #hmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxstart
                        fake_imgs, layers_output, _, _ = netG(
                            noise_j, t_embeddings_i)
                        if len(layers_output) != len(lamdas):
                            print("please check lamdas length")
                        #hmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxhmxend
                        # filenames_number ='_sample_%2.2d'%(j)
                        # filenames_new = []
                        # filenames_new.append(filenames[-1]+filenames_number)
                        # filenames_new = tuple(filenames_new)

                        if cfg.TEST.B_EXAMPLE:
                            # fake_img_list.append(fake_imgs[0].data.cpu())
                            # fake_img_list.append(fake_imgs[1].data.cpu())
                            fake_img_list.append(fake_imgs[2].data.cpu())
                        else:
                            self.save_singleimages(fake_imgs[-1], filenames, j,
                                                   save_dir, split_dir, i, 256)
                        # self.save_singleimages(fake_imgs[-2], filenames,
                        #                        save_dir, split_dir, i, 128)
                        # self.save_singleimages(fake_imgs[-3], filenames,
                        #                        save_dir, split_dir, i, 64)
                    # break
                if cfg.TEST.B_EXAMPLE:
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 64)
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 128)
                    self.save_superimages(fake_img_list, filenames, save_dir,
                                          split_dir, 256)
コード例 #21
0
    def build_models(self):
        # ###################encoders######################################## #
        if cfg.TRAIN.NET_E == '':
            print('Error: no pretrained text-image encoders')
            return

        image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
        img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
                                                   'image_encoder')
        state_dict = \
            torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
        image_encoder.load_state_dict(state_dict)
        for p in image_encoder.parameters():
            p.requires_grad = False
        print('Load image encoder from:', img_encoder_path)
        image_encoder.eval()

        text_encoder = \
            RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
        state_dict = \
            torch.load(cfg.TRAIN.NET_E,
                       map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        for p in text_encoder.parameters():
            p.requires_grad = False
        print('Load text encoder from:', cfg.TRAIN.NET_E)
        text_encoder.eval()

        # #######################generator and discriminators############## #
        netG = G_NET(len(self.cats_index_dict))
        netsPatD, netsShpD = [], []
        if cfg.TREE.BRANCH_NUM > 0:
            netsPatD.append(PAT_D_NET64())
            netsShpD.append(SHP_D_NET64(len(self.cats_index_dict)))
        if cfg.TREE.BRANCH_NUM > 1:
            netsPatD.append(PAT_D_NET128())
            netsShpD.append(SHP_D_NET128(len(self.cats_index_dict)))
        if cfg.TREE.BRANCH_NUM > 2:
            netsPatD.append(PAT_D_NET256())
            netsShpD.append(SHP_D_NET256(len(self.cats_index_dict)))

        netObjSSD = OBJ_SS_D_NET(len(self.cats_index_dict))
        netObjLSD = OBJ_LS_D_NET(len(self.cats_index_dict))

        netG.apply(weights_init)
        netObjSSD.apply(weights_init)
        netObjLSD.apply(weights_init)
        for i in range(len(netsPatD)):
            netsPatD[i].apply(weights_init)
            netsShpD[i].apply(weights_init)
        print('# of netsPatD', len(netsPatD))
        # ########################################################### #
        if cfg.CUDA:
            text_encoder = text_encoder.cuda()
            image_encoder = image_encoder.cuda()
            netG.cuda()
            netObjSSD.cuda()
            netObjLSD.cuda()
            for i in range(len(netsPatD)):
                netsPatD[i].cuda()
                netsShpD[i].cuda()

            if len(cfg.GPU_IDS) > 1:
                text_encoder = nn.DataParallel(text_encoder)
                text_encoder.to(self.device)
                image_encoder = nn.DataParallel(image_encoder)
                image_encoder.to(self.device)
                netG = nn.DataParallel(netG)
                netG.to(self.device)
                netObjSSD = nn.DataParallel(netObjSSD)
                netObjSSD.to(self.device)
                netObjLSD = nn.DataParallel(netObjLSD)
                netObjLSD.to(self.device)
                for i in range(len(netsPatD)):
                    netsPatD[i] = nn.DataParallel(netsPatD[i])
                    netsPatD[i].to(self.device)

                    netsShpD[i] = nn.DataParallel(netsShpD[i])
                    netsShpD[i].to(self.device)
        #
        epoch = 0
        if cfg.TRAIN.NET_G != '':
            state_dict = \
                torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load G from: ', cfg.TRAIN.NET_G)
            istart = cfg.TRAIN.NET_G.rfind('_') + 1
            iend = cfg.TRAIN.NET_G.rfind('.')
            epoch = cfg.TRAIN.NET_G[istart:iend]
            epoch = int(epoch) + 1

            Gname = cfg.TRAIN.NET_G
            for i in range(len(netsPatD)):
                s_tmp = Gname[:Gname.rfind('/')]

                Dname = '%s/netPatD%d.pth' % (s_tmp, i)
                print('Load PatD from: ', Dname)
                state_dict = \
                    torch.load(Dname, map_location=lambda storage, loc: storage)
                netsPatD[i].load_state_dict(state_dict)

                Dname = '%s/netShpD%d.pth' % (s_tmp, i)
                print('Load ShpD from: ', Dname)
                state_dict = \
                    torch.load(Dname, map_location=lambda storage, loc: storage)
                netsShpD[i].load_state_dict(state_dict)

            s_tmp = Gname[:Gname.rfind('/')]
            Dname = '%s/netObjSSD.pth' % (s_tmp)
            print('Load ObjSSD from: ', Dname)
            state_dict = \
                torch.load(Dname, map_location=lambda storage, loc: storage)
            netObjSSD.load_state_dict(state_dict)

            s_tmp = Gname[:Gname.rfind('/')]
            Dname = '%s/netObjLSD.pth' % (s_tmp)
            print('Load ObjLSD from: ', Dname)
            state_dict = \
                torch.load(Dname, map_location=lambda storage, loc: storage)
            netObjLSD.load_state_dict(state_dict)

        return [
            text_encoder, image_encoder, netG, netsPatD, netsShpD, netObjSSD,
            netObjLSD, epoch
        ]
コード例 #22
0
    def sample(self, split_dir, num_samples=25, draw_bbox=False):
        from PIL import Image, ImageDraw, ImageFont
        import cPickle as pickle
        import torchvision
        import torchvision.utils as vutils

        if cfg.TRAIN.NET_G == '':
            print('Error: the path for model NET_G is not found!')
        else:
            if split_dir == 'test':
                split_dir = 'valid'
            # Build and load the generator
            text_encoder = RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
            state_dict = \
                torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
            text_encoder.load_state_dict(state_dict)
            print('Load text encoder from:', cfg.TRAIN.NET_E)
            text_encoder = text_encoder.cuda()
            text_encoder.eval()

            batch_size = cfg.TRAIN.BATCH_SIZE
            nz = cfg.GAN.Z_DIM

            model_dir = cfg.TRAIN.NET_G
            state_dict = torch.load(model_dir, map_location=lambda storage, loc: storage)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            netG = G_NET()
            print('Load G from: ', model_dir)
            netG.apply(weights_init)

            netG.load_state_dict(state_dict["netG"])
            netG.cuda()
            netG.eval()

            # the path to save generated images
            s_tmp = model_dir[:model_dir.rfind('.pth')]
            save_dir = '%s_%s' % (s_tmp, split_dir)
            mkdir_p(save_dir)
            #######################################
            noise = Variable(torch.FloatTensor(9, nz))

            imsize = 256

            for step, data in enumerate(self.data_loader, 0):
                if step >= num_samples:
                    break

                imgs, captions, cap_lens, class_ids, keys, transformation_matrices, label_one_hot, bbox = \
                    prepare_data(data, eval=True)
                transf_matrices_inv = transformation_matrices[1][0].unsqueeze(0)
                label_one_hot = label_one_hot[0].unsqueeze(0)

                img = imgs[-1][0]
                val_image = img.view(1, 3, imsize, imsize)

                hidden = text_encoder.init_hidden(batch_size)
                # words_embs: batch_size x nef x seq_len
                # sent_emb: batch_size x nef
                words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
                words_embs, sent_emb = words_embs[0].unsqueeze(0).detach(), sent_emb[0].unsqueeze(0).detach()
                words_embs = words_embs.repeat(9, 1, 1)
                sent_emb = sent_emb.repeat(9, 1)
                mask = (captions == 0)
                mask = mask[0].unsqueeze(0)
                num_words = words_embs.size(2)
                if mask.size(1) > num_words:
                    mask = mask[:, :num_words]
                mask = mask.repeat(9, 1)
                transf_matrices_inv = transf_matrices_inv.repeat(9, 1, 1, 1)
                label_one_hot = label_one_hot.repeat(9, 1, 1)

                #######################################################
                # (2) Generate fake images
                ######################################################
                noise.data.normal_(0, 1)
                inputs = (noise, sent_emb, words_embs, mask, transf_matrices_inv, label_one_hot)
                with torch.no_grad():
                    fake_imgs, _, mu, logvar = nn.parallel.data_parallel(netG, inputs, self.gpus)

                data_img = torch.FloatTensor(10, 3, imsize, imsize).fill_(0)
                data_img[0] = val_image
                data_img[1:10] = fake_imgs[-1]

                if draw_bbox:
                    for idx in range(3):
                        x, y, w, h = tuple([int(imsize*x) for x in bbox[0, idx]])
                        w = imsize-1 if w > imsize-1 else w
                        h = imsize-1 if h > imsize-1 else h
                        if x <= -1:
                            break
                        data_img[:10, :, y, x:x + w] = 1
                        data_img[:10, :, y:y + h, x] = 1
                        data_img[:10, :, y+h, x:x + w] = 1
                        data_img[:10, :, y:y + h, x + w] = 1

                # get caption
                cap = captions[0].data.cpu().numpy()
                sentence = ""
                for j in range(len(cap)):
                    if cap[j] == 0:
                        break
                    word = self.ixtoword[cap[j]].encode('ascii', 'ignore').decode('ascii')
                    sentence += word + " "
                sentence = sentence[:-1]
                vutils.save_image(data_img, '{}/{}_{}.png'.format(save_dir, sentence, step), normalize=True, nrow=10)
            print("Saved {} files to {}".format(step, save_dir))
コード例 #23
0
ファイル: trainer.py プロジェクト: zxs789/Obj-GAN
    def sampling(self, split_dir):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            if split_dir == 'test':
                split_dir = 'valid'
            # Build and load the generator
            netG = G_NET(len(self.cats_index_dict))
            netG.apply(weights_init)
            netG.eval()

            if cfg.CUDA:
                netG.cuda()

            if len(cfg.GPU_IDS) > 1:
                netG = nn.DataParallel(netG)
                netG.to(self.device)

            batch_size = self.batch_size
            nz = cfg.GAN.Z_DIM
            noise = Variable(
                torch.FloatTensor(batch_size, cfg.ROI.BOXES_NUM,
                                  len(self.cats_index_dict) * 4))
            noise = noise.cuda()

            model_dir = cfg.TRAIN.NET_G
            state_dict = \
                torch.load(model_dir, map_location=lambda storage, loc: storage)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            netG.load_state_dict(state_dict)
            print('Load G from: ', model_dir)

            # the path to save generated images
            s_tmp = model_dir[:model_dir.rfind('.pth')]
            save_dir = '%s/%s' % (s_tmp, split_dir)
            mkdir_p(save_dir)

            cnt = 0

            for _ in range(1):  # (cfg.TEXT.CAPTIONS_PER_IMAGE):
                for step, data in enumerate(self.data_loader, 0):
                    cnt += batch_size
                    if step % 100 == 0:
                        print('step: ', step)
                    # if step > 50:
                    #     break
                    imgs, pooled_hmaps, hmaps, bbox_maps_fwd, bbox_maps_bwd, bbox_fmaps, \
                        rois, fm_rois, num_rois, class_ids, keys = prepare_data(data)
                    num_rois = num_rois.data.cpu().numpy()

                    cats_list = []
                    for batch_index in range(self.batch_size):
                        cats = []
                        for roi_index in range(num_rois[batch_index]):
                            rela_cat_id = int(rois[batch_index, roi_index, 4])
                            abs_cat_id = self.cats_dict[rela_cat_id][0]
                            cat = self.ixtoword[abs_cat_id].encode(
                                'ascii', 'ignore').decode('ascii')
                            cats.append(cat)
                        cats_list.append(cats)

                    #######################################################
                    # (2) Generate fake images
                    ######################################################
                    max_num_roi = max(num_rois)
                    noise.data.normal_(0, 1)
                    fake_hmaps = netG(noise[:, :max_num_roi], bbox_maps_fwd,
                                      bbox_maps_bwd, bbox_fmaps)
                    fake_hmaps = fake_hmaps.repeat(1, 1, 3, 1, 1)
                    for j in range(batch_size):
                        s_tmp = '%s/single/%s' % (save_dir, keys[j])
                        folder = s_tmp[:s_tmp.rfind('/')]
                        if not os.path.isdir(folder):
                            print('Make a new folder: ', folder)
                            mkdir_p(folder)
                        k = 0
                        # for k in range(len(fake_imgs)):
                        im = fake_hmaps[j][k].data.cpu().numpy()

                        minV = im.min()
                        maxV = im.max()
                        im = (im - minV) / (maxV - minV)
                        im *= 255
                        im = im.astype(np.uint8)
                        im = np.transpose(im, (1, 2, 0))
                        im = Image.fromarray(im)

                        cat = cats_list[j][k]
                        fullpath = '{0}_{1}.png'.format(s_tmp, cat)
                        im.save(fullpath)
コード例 #24
0
nz = 100
n_samples = 10

fid_model = FID_INCEPTION()
fid_model.cuda()
fid_model.eval()

inception_model = INCEPTION_V3()
inception_model.cuda()
inception_model.eval()

G_NET_Path = 'C:\\Users\\alper\\PycharmProjects\\MSGAN\\StackGAN++-Mode-Seeking\\models\\ours_new.pth'
netG = G_NET()
netG.apply(weights_init)
torch.cuda.set_device(0)
netG = netG.cuda()
netG = torch.nn.DataParallel(netG, device_ids=[0])
state_dict = \
    torch.load(G_NET_Path,
               map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)


noise = Variable(torch.FloatTensor(n_samples, nz))
noise = noise.cuda()
netG.eval()

def generate_fake_images():
    emb = dataset.embeddings
    emb = np.reshape(emb, (emb.shape[0]*emb.shape[1], emb.shape[2]))
    np.random.shuffle(emb)
コード例 #25
0
class Generator:
    def __init__(self, caption_file, saveable, cuda=False, profile=False):
        # flags
        self.cuda = cuda
        self.profile = profile

        if self.profile:
            print('Initializing Generator...')
            print('cuda={}\nprofile={}'.format(self.cuda, self.profile))

        # load caption indices
        x = pickle.load(open(caption_file, 'rb'))
        self.ixtoword = x[2]
        self.wordtoix = x[3]
        del x

        # load text encoder
        self.text_encoder = RNN_ENCODER(len(self.wordtoix), nhidden=cfg.TEXT.EMBEDDING_DIM)
        state_dict = torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
        self.text_encoder.load_state_dict(state_dict)
        if self.cuda:
            self.text_encoder.cuda()
            
        self.text_encoder.eval()

        # load generative model
        self.netG = G_NET()
        state_dict = torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
        self.netG.load_state_dict(state_dict)
        if self.cuda:
            self.netG.cuda()
            
        self.netG.eval()

        # saveable items -> push to storage
        self.saveable = saveable

    def vectorize_caption(self, caption, copies):
        # create caption vector
        tokens = caption.split(' ')
        cap_v = []
        for t in tokens:
            t = t.strip().encode('ascii', 'ignore').decode('ascii')
            if len(t) > 0 and t in self.wordtoix:
                cap_v.append(self.wordtoix[t])

        # expected state for single generation
        captions = np.zeros((copies, len(cap_v)))
        for i in range(copies):
            captions[i,:] = np.array(cap_v)
        cap_lens = np.zeros(copies) + len(cap_v)

        return captions.astype(int), cap_lens.astype(int), len(self.wordtoix)

    def generate(self, caption, copies=2):
        # load word vector
        captions, cap_lens, n_words = self.vectorize_caption(caption, copies)

        # only one to generate
        batch_size = captions.shape[0]

        nz = cfg.GAN.Z_DIM
        captions = Variable(torch.from_numpy(captions), volatile=True)
        cap_lens = Variable(torch.from_numpy(cap_lens), volatile=True)
        noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)

        if self.cuda:
            captions = captions.cuda()
            cap_lens = cap_lens.cuda()
            noise = noise.cuda()

        #######################################################
        # (1) Extract text embeddings
        #######################################################
        hidden = self.text_encoder.init_hidden(batch_size)
        words_embs, sent_emb = self.text_encoder(captions, cap_lens, hidden)
        mask = (captions == 0)

        #######################################################
        # (2) Generate fake images
        #######################################################
        noise.data.normal_(0, 1)
        fake_imgs, attention_maps, _, _ = self.netG(noise, sent_emb, words_embs, mask)

        # G attention
        cap_lens_np = cap_lens.cpu().data.numpy()

        # prefix for partitioning images
        prefix = datetime.now().strftime('%Y/%B/%d/%H_%M_%S_%f')
        urls = []
        # only look at first one
        for j in range(batch_size):
            for k in range(len(fake_imgs)):
                im = fake_imgs[k][j].data.cpu().numpy()
                im = (im + 1.0) * 127.5
                im = im.astype(np.uint8)
                im = np.transpose(im, (1, 2, 0))

                # save using saveable
                birdy = 'bird_g{}'.format(k)
                if copies > 2:
                    item = self.saveable.save('{}/{}'.format(prefix, j), birdy, im)
                else:
                    item = self.saveable.save(prefix, birdy, im)

                urls.append(item)

            if copies == 2:
                for k in range(len(attention_maps)):
                    if len(fake_imgs) > 1:
                        im = fake_imgs[k + 1].detach().cpu()
                    else:
                        im = fake_imgs[0].detach().cpu()
                            
                    attn_maps = attention_maps[k]
                    att_sze = attn_maps.size(2)

                    img_set, sentences = \
                        build_super_images2(im[j].unsqueeze(0),
                                            captions[j].unsqueeze(0),
                                            [cap_lens_np[j]], self.ixtoword,
                                            [attn_maps[j]], att_sze)

                    if img_set is not None:
                        attnmap = 'attmaps_a{}'.format(k)
                        item = self.saveable.save(prefix, attnmap, img_set)
                        urls.append(item)
            if copies == 2:
                break

        return urls
コード例 #26
0
    def sampling(self):
        if self.args.netG == '':
            print('Error: the path for models is not found!')
        else:
            data_dir = cfg.DATA_DIR
            if self.args.split == "test_unseen":
                filepath = os.path.join(data_dir,
                                        "test_unseen/class_data.pickle")
            else:  #test_seen
                filepath = os.path.join(data_dir,
                                        "test_seen/class_data.pickle")
            if os.path.isfile(filepath):
                with open(filepath, "rb") as f:
                    data_dic = pkl.load(f)
            class_names = data_dic['classes']
            class_ids = data_dic['class_info']

            att_dir = os.path.join(data_dir, "CUB_200_2011/attributes")
            att_np = np.zeros((312, 200))  #for CUB
            with open(att_dir + "/class_attribute_labels_continuous.txt",
                      "r") as f:
                for ind, line in enumerate(f.readlines()):
                    line = line.strip("\n")
                    line = list(map(float, line.split()))
                    att_np[:, ind] = line

            if self.args.kl_loss:
                netG = G_NET()
            else:
                netG = G_NET_not_CA()
            test_model = "netG_epoch_600.pth"
            model_path = os.path.join(self.args.netG, "Model", test_model)  ##
            state_dic = torch.load(model_path,
                                   map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dic)
            netG.cuda()
            netG.eval()

            noise = torch.FloatTensor(self.batch_size, cfg.GAN.Z_DIM)

            for class_name, class_id in zip(class_names, class_ids):
                print("now generating, ", class_name)
                class_dir = os.path.join(self.args.netG, 'valid',
                                         test_model[:test_model.rfind(".")],
                                         self.args.split, class_name)
                atts = att_np[:, class_id - 1]
                atts = np.expand_dims(atts, axis=0)
                atts = atts.repeat(self.batch_size, axis=0)
                assert atts.shape == (self.batch_size, 312)

                if cfg.CUDA:
                    noise = noise.cuda()
                    atts = torch.cuda.FloatTensor(atts)
                else:
                    atts = torch.FloatTensor(atts)

                for i in range(self.sample_num):
                    noise.normal_(0, 1)
                    if self.args.kl_loss:
                        fake_imgs, _, _ = nn.parallel.data_parallel(
                            netG, (noise, atts), self.gpus)
                    else:
                        fake_imgs = nn.parallel.data_parallel(
                            netG, (noise, atts), self.gpus)
                    for stage in range(len(fake_imgs)):
                        for num, im in enumerate(fake_imgs[stage]):
                            im = im.detach().cpu()
                            im = im.add_(1).div_(2).mul_(255)
                            im = im.numpy().astype(np.uint8)
                            im = np.transpose(im, (1, 2, 0))
                            im = Image.fromarray(im)
                            stage_dir = os.path.join(class_dir,
                                                     "stage_%d" % stage)
                            mkdir_p(stage_dir)
                            img_path = os.path.join(stage_dir,
                                                    "single_%d.png" % num)
                            im.save(img_path)
                        for j in range(int(self.batch_size /
                                           20)):  ## cfg.batch_size==100
                            one_set = [
                                fake_imgs[0][j * 20:(j + 1) * 20],
                                fake_imgs[1][j * 20:(j + 1) * 20],
                                fake_imgs[2][j * 20:(j + 1) * 20]
                            ]
                            img_set = build_images(one_set)
                            img_set = Image.fromarray(img_set)
                            super_dir = os.path.join(class_dir, "super")
                            mkdir_p(super_dir)
                            img_path = os.path.join(super_dir,
                                                    "super_%d.png" % j)
                            img_set.save(img_path)
コード例 #27
0
    def build_models(self):
        # ############################## encoders ############################# #
        text_encoder = RNN_ENCODER(self.n_words,
                                   nhidden=cfg.TEXT.EMBEDDING_DIM)
        state_dict = \
            torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        print('Load text encoder from:', cfg.TRAIN.NET_E)
        text_encoder.eval()

        image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
        img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder',
                                                   'image_encoder')
        state_dict = \
            torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
        image_encoder.load_state_dict(state_dict)
        for p in image_encoder.parameters():
            p.requires_grad = False
        print('Load image encoder from:', img_encoder_path)
        image_encoder.eval()

        # ########### image generator and (potential) shape generator ########## #
        netG = G_NET(len(self.cats_index_dict))
        netG.apply(weights_init)
        netG.eval()
        netShpG = None
        if cfg.TEST.USE_GT_BOX_SEG > 0:
            netShpG = SHP_G_NET(len(self.cats_index_dict))
            netShpG.apply(weights_init)
            netShpG.eval()

        # ################### parallization and initialization ################## #
        if cfg.CUDA:
            text_encoder.cuda()
            image_encoder.cuda()
            netG.cuda()
            if cfg.TEST.USE_GT_BOX_SEG > 0:
                netShpG.cuda()

            if len(cfg.GPU_IDS) > 1:
                text_encoder = nn.DataParallel(text_encoder)
                text_encoder.to(self.device)
                image_encoder = nn.DataParallel(image_encoder)
                image_encoder.to(self.device)
                netG = nn.DataParallel(netG)
                netG.to(self.device)

            if cfg.TEST.USE_GT_BOX_SEG > 0:
                netShpG = nn.DataParallel(netShpG)
                netShpG.to(self.device)

        state_dict = torch.load(cfg.TRAIN.NET_G,
                                map_location=lambda storage, loc: storage)
        netG.load_state_dict(state_dict)
        print('Load G from: ', cfg.TRAIN.NET_G)

        if cfg.TEST.USE_GT_BOX_SEG > 0:
            state_dict = torch.load(cfg.TEST.NET_SHP_G,
                                    map_location=lambda storage, loc: storage)
            netShpG.load_state_dict(state_dict)
            print('Load Shape G from: ', cfg.TEST.NET_SHP_G)

        return [text_encoder, image_encoder, netG, netShpG]
コード例 #28
0
    def evaluate(self, split_dir):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            # Build and load the generator
            if split_dir == 'test':
                split_dir = 'valid'
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            print(netG)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            netE = load_embedding_model(self.data_loader.dataset.dictionary)
            print(netE)

            nz = cfg.GAN.Z_DIM
            sample_size = cfg.TEST.NUM_IMAGES
            noise = Variable(torch.FloatTensor(sample_size, nz))
            if cfg.CUDA:
                netG.cuda()
                netE.cuda()
                noise = noise.cuda()

            # switch to evaluate mode
            netG.eval()
            count = 0
            output_dir = os.path.join(cfg.OUTPUT_DIR, cfg.EXPERIMENT_NAME)
            for step, data in enumerate(
                    tqdm(self.data_loader, desc='evaluate'), 0):
                imgs, txt_ids, txts = data

                if cfg.CUDA:
                    txt_ids = Variable(txt_ids).cuda()
                else:
                    txt_ids = Variable(txt_ids)

                txts_embeddings = netE(txt_ids)

                batch_size = imgs[0].size(0)

                imgs64, imgs128, imgs256 = [], [], []
                for i in range(0, batch_size):
                    noise.data.normal_(0, 1)
                    txt_embedding = txts_embeddings[i].repeat(sample_size, 1)

                    fake_imgs, _, _ = netG(noise, txt_embedding)

                    imgs64.append(normalize_(fake_imgs[0]))
                    imgs128.append(normalize_(fake_imgs[1]))
                    imgs256.append(normalize_(fake_imgs[2]))

                save_images_with_text(imgs64, imgs128, imgs256, imgs, txts,
                                      batch_size, cfg.TEXT.MAX_LEN, count,
                                      output_dir)

                count = count + batch_size + 1
コード例 #29
0
ファイル: test.py プロジェクト: pan-x-c/Wenhua_Text_to_Image
torch.cuda.set_device(gpus[0])
cudnn.benchmark = True
batch_size = 2

netG = G_NET()
netG.apply(weights_init)
netG = torch.nn.DataParallel(netG, device_ids=gpus)
state_dict = \
    torch.load(cfg.TRAIN.NET_G,
               map_location=lambda storage, loc: storage)
netG.load_state_dict(state_dict)

nz = cfg.GAN.Z_DIM
noise = Variable(torch.FloatTensor(batch_size, nz))

netG.cuda()
netG.eval()
noise = noise.cuda()

t_embeddings = load_lua(txt_dir)
t_embeddings = t_embeddings.view(-1,1024)

t_embeddings = torch.cat((t_embeddings,t_embeddings), 0)
t_embeddings = Variable(t_embeddings).cuda()

embedding_dim = t_embeddings.size(1)

noise.data.resize_(batch_size, nz)
noise.data.normal_(0, 1)
images, _, _ = netG(noise, t_embeddings)
save_singleimages(images[-1], save_dir)
コード例 #30
0
    def sample_images(self):
        sample_size = 24
        save_dir = '../sample_images/'
        save_final = '../sample_finals/'

        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        if not os.path.exists(save_final):
            os.makedirs(save_final)

        random.seed(datetime.now())
        depth = cfg.TEST_DEPTH
        res = 32 * 2**depth
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for model not found!')
        else:
            # Build and load the generator
            netG = G_NET(depth)
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            model_dict = netG.state_dict()

            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)

            state_dict = {
                k: v
                for k, v in state_dict.items() if k in model_dict
            }

            model_dict.update(state_dict)
            netG.load_state_dict(model_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # Uncomment this to print Generator layers
            # print(netG)

            nz = cfg.GAN.Z_DIM
            noise = torch.FloatTensor(1, nz)
            # noise.data.normal_(0, 1)
            # noise = noise.repeat(1, 1)

            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            netG.eval()

            for i in tqdm(range(sample_size)):
                noise.data.normal_(0, 1)
                bg_code = torch.zeros([1, cfg.FINE_GRAINED_CATEGORIES]).cuda()
                p_code = torch.zeros([1, cfg.SUPER_CATEGORIES]).cuda()
                c_code = torch.zeros([1, cfg.FINE_GRAINED_CATEGORIES]).cuda()
                b = random.randint(0, cfg.FINE_GRAINED_CATEGORIES - 1)
                p = random.randint(0, cfg.SUPER_CATEGORIES - 1)
                c = random.randint(0, cfg.FINE_GRAINED_CATEGORIES - 1)
                bg_code[0][b] = 1
                p_code[0][p] = 1
                c_code[0][c] = 1

                fake_imgs, fg_imgs, mk_imgs, fgmk_imgs = netG(
                    noise, c_code, 1, p_code,
                    bg_code)  # Forward pass through the generator

                self.save_image(fake_imgs[3 * depth + 0][0], save_dir,
                                '%d_bg' % i)
                self.save_image(fake_imgs[3 * depth + 1][0], save_dir,
                                '%d_pf' % i)
                self.save_image(fake_imgs[3 * depth + 2][0], save_dir,
                                '%d_cf' % i)
                self.save_image(fake_imgs[3 * depth + 2][0], save_final,
                                '%d' % i)
                # self.save_image(fg_imgs[2 * depth + 0][0], save_dir, 'parent_foreground')
                # self.save_image(fg_imgs[2 * depth + 1][0], save_dir, 'child_foreground')
                self.save_image(mk_imgs[2 * depth + 0][0], save_dir,
                                '%d_pmk' % i)
                self.save_image(mk_imgs[2 * depth + 1][0], save_dir,
                                '%d_cmk' % i)
コード例 #31
0
ファイル: trainer.py プロジェクト: gchochla/StackGAN-v2
    def evaluate(self,
                 split_dir,
                 n_samples=4,
                 extractor='googlenet',
                 save_dir=None):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            # Build and load the generator
            if split_dir == 'test':
                split_dir = 'valid'
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            mapper = EXTRACTOR_MAPPING[extractor]()
            mapper = torch.nn.DataParallel(mapper, device_ids=self.gpus)
            set_parameter_requires_grad(netG, False)
            set_parameter_requires_grad(mapper, False)

            print(netG)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            if save_dir is None:
                # the path to save generated images
                s_tmp = cfg.TRAIN.NET_G
                istart = s_tmp.rfind('_') + 1
                iend = s_tmp.rfind('.')
                iteration = int(s_tmp[istart:iend])
                s_tmp = s_tmp[:s_tmp.rfind('/')]
                save_dir = '%s/iteration%d' % (s_tmp, iteration)

            nz = cfg.GAN.Z_DIM

            if cfg.CUDA:
                netG.cuda()
                mapper.cuda()

            # switch to evaluate mode
            netG.eval()
            mapper.eval()

            synthetic_ds = SyntheticDataset(save_dir)

            for class_embeddings, synthetic_id in self.data_loader.dataset.embeddings_by_class(
            ):

                if cfg.CUDA:
                    class_embeddings = class_embeddings.cuda()

                class_embeddings = class_embeddings.mean(
                    dim=1)  # mean of 10 captions per image
                for i in range(class_embeddings.size(0)):
                    image_embeddings = class_embeddings[i].repeat(n_samples, 1)
                    noise = torch.randn(n_samples, nz)
                    if cfg.CUDA:
                        noise = noise.cuda()

                    imgs, _, _ = netG(noise, image_embeddings)
                    imgs = imgs[-1]
                    samples = mapper(imgs)

                    synthetic_ds.save_pairs(samples, synthetic_id)
コード例 #32
0
    def build_models(self):
        # ###################encoders######################################## #
        if cfg.TRAIN.NET_E == '':
            print('Error: no pretrained text-image encoders')
            return

        image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
        img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
        state_dict = \
            torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
        image_encoder.load_state_dict(state_dict)
        for p in image_encoder.parameters():
            p.requires_grad = False
        print('Load image encoder from:', img_encoder_path)
        image_encoder.eval()

        text_encoder = \
            RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
        state_dict = \
            torch.load(cfg.TRAIN.NET_E,
                       map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        for p in text_encoder.parameters():
            p.requires_grad = False
        print('Load text encoder from:', cfg.TRAIN.NET_E)
        text_encoder.eval()

        # #######################generator and discriminators############## #
        netsD = []
        from model import D_NET64, D_NET128, D_NET256
        netG = G_NET()
        if cfg.TREE.BRANCH_NUM > 0:
            netsD.append(D_NET64())
        if cfg.TREE.BRANCH_NUM > 1:
            netsD.append(D_NET128())
        if cfg.TREE.BRANCH_NUM > 2:
            netsD.append(D_NET256())

        netG.apply(weights_init)
        # print(netG)
        for i in range(len(netsD)):
            netsD[i].apply(weights_init)
            # print(netsD[i])
        print('# of netsD', len(netsD))
        epoch = 0

        if self.resume:
            checkpoint_list = sorted([ckpt for ckpt in glob.glob(self.model_dir + "/" + '*.pth')])
            latest_checkpoint = checkpoint_list[-1]
            state_dict = torch.load(latest_checkpoint, map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict["netG"])
            for i in range(len(netsD)):
                netsD[i].load_state_dict(state_dict["netD"][i])
            epoch = int(latest_checkpoint[-8:-4]) + 1
            print("Resuming training from checkpoint {} at epoch {}.".format(latest_checkpoint, epoch))

        #
        if cfg.TRAIN.NET_G != '':
            state_dict = \
                torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load G from: ', cfg.TRAIN.NET_G)
            istart = cfg.TRAIN.NET_G.rfind('_') + 1
            iend = cfg.TRAIN.NET_G.rfind('.')
            epoch = cfg.TRAIN.NET_G[istart:iend]
            epoch = int(epoch) + 1
            if cfg.TRAIN.B_NET_D:
                Gname = cfg.TRAIN.NET_G
                for i in range(len(netsD)):
                    s_tmp = Gname[:Gname.rfind('/')]
                    Dname = '%s/netD%d.pth' % (s_tmp, i)
                    print('Load D from: ', Dname)
                    state_dict = \
                        torch.load(Dname, map_location=lambda storage, loc: storage)
                    netsD[i].load_state_dict(state_dict)
        # ########################################################### #
        if cfg.CUDA:
            text_encoder = text_encoder.cuda()
            image_encoder = image_encoder.cuda()
            netG.cuda()
            for i in range(len(netsD)):
                netsD[i].cuda()
        return [text_encoder, image_encoder, netG, netsD, epoch]
コード例 #33
0
ファイル: trainer.py プロジェクト: tensoralex/StackGAN-v2
    def evaluate(self, split_dir):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            # Build and load the generator
            if split_dir == 'test':
                split_dir = 'valid'
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            print(netG)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # the path to save generated images
            s_tmp = cfg.TRAIN.NET_G
            istart = s_tmp.rfind('_') + 1
            iend = s_tmp.rfind('.')
            iteration = int(s_tmp[istart:iend])
            s_tmp = s_tmp[:s_tmp.rfind('/')]
            save_dir = '%s/iteration%d' % (s_tmp, iteration)

            nz = cfg.GAN.Z_DIM
            noise = Variable(torch.FloatTensor(self.batch_size, nz))
            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            # switch to evaluate mode
            netG.eval()
            for step, data in enumerate(self.data_loader, 0):
                imgs, t_embeddings, filenames = data
                if cfg.CUDA:
                    t_embeddings = Variable(t_embeddings).cuda()
                else:
                    t_embeddings = Variable(t_embeddings)
                # print(t_embeddings[:, 0, :], t_embeddings.size(1))

                embedding_dim = t_embeddings.size(1)
                batch_size = imgs[0].size(0)
                noise.data.resize_(batch_size, nz)
                noise.data.normal_(0, 1)

                fake_img_list = []
                for i in range(embedding_dim):
                    fake_imgs, _, _ = netG(noise, t_embeddings[:, i, :])
                    if cfg.TEST.B_EXAMPLE:
                        # fake_img_list.append(fake_imgs[0].data.cpu())
                        # fake_img_list.append(fake_imgs[1].data.cpu())
                        fake_img_list.append(fake_imgs[2].data.cpu())
                    else:
                        self.save_singleimages(fake_imgs[-1], filenames,
                                               save_dir, split_dir, i, 256)
                        # self.save_singleimages(fake_imgs[-2], filenames,
                        #                        save_dir, split_dir, i, 128)
                        # self.save_singleimages(fake_imgs[-3], filenames,
                        #                        save_dir, split_dir, i, 64)
                    # break
                if cfg.TEST.B_EXAMPLE:
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 64)
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 128)
                    self.save_superimages(fake_img_list, filenames,
                                          save_dir, split_dir, 256)