Exemple #1
0
netd = NetD(nc=1)
neta = NetA(nc=1)

device = th.device("cuda:0")
netg = netg.to(device)
netd = netd.to(device)
neta = neta.to(device)
fineSize = 64

checkpoint = '/home/mg/code/my_GAN_dataSet/snapshots/snapshot_449.t7'
checkpoint = th.load(checkpoint)
neta.load_state_dict(checkpoint['netA'])
netg.load_state_dict(checkpoint['netG'])
netd.load_state_dict(checkpoint['netD'])
neta.eval()
netg.eval()
netd.eval()

angles = [
    '000', '018', '036', '054', '072', '090', '108', '126', '144', '162', '180'
]

for cond in ['nm-01', 'nm-02', 'nm-03', 'nm-04', 'cl-01', 'cl-02']:
    dataset = CASIABDatasetGenerate(
        data_dir='/home/mg/code/data/GEI_CASIA_B/gei/', cond=cond)
    for i in range(1, 125):
        ass_label, img = dataset.getbatch(i, 11)
        img = img.to(device).to(th.float32)

        with th.no_grad():
            fake = netg(img)
Exemple #2
0
def main(args):
    # manualSeed to control the noise
    manualSeed = 100
    random.seed(manualSeed)
    np.random.seed(manualSeed)
    torch.manual_seed(manualSeed)

    with open(args.json_file, 'r') as f:
        dataset_json = json.load(f)

    # load rnn encoder
    text_encoder = RNN_ENCODER(dataset_json['n_words'], nhidden=dataset_json['text_embed_dim'])
    text_encoder_dir = args.rnn_encoder
    state_dict = torch.load(text_encoder_dir, map_location=lambda storage, loc: storage)
    text_encoder.load_state_dict(state_dict)

    # load netG
    state_dict = torch.load(args.model_path, map_location=torch.device('cpu'))
    # netG = NetG(int(dataset_json['n_channels']), int(dataset_json['cond_dim']))
    netG = NetG(64, int(dataset_json['cond_dim']))
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:]  # remove `module.`nvidia
        new_state_dict[name] = v
    model_dict = netG.state_dict()
    pretrained_dict = {k: v for k, v in new_state_dict.items() if k in model_dict}
    model_dict.update(pretrained_dict)
    netG.load_state_dict(model_dict)

    # use gpu or not, change model to evaluation mode
    if args.use_gpu:
        text_encoder.cuda()
        netG.cuda()
        caption_idx.cuda()
        caption_len.cuda()
        noise.cuda()

    text_encoder.eval()
    netG.eval()

    # generate noise
    num_noise = 100
    noise = torch.FloatTensor(num_noise, 100)

    # cub bird captions
    # caption = 'this small bird has a light yellow breast and brown wings'
    # caption = 'this small bird has a short beak a light gray breast a darker gray crown and black wing tips'
    # caption = 'this small bird has wings that are gray and has a white belly'
    # caption = 'this bird has a yellow throat belly abdomen and sides with lots of brown streaks on them'
    # caption = 'this little bird has a yellow belly and breast with a gray wing with white wingbars'
    # caption = 'this bird has a white belly and breast wit ha blue crown and nape'
    # caption = 'a bird with brown and black wings red crown and throat and the bill is short and pointed'
    # caption = 'this small bird has a yellow crown and a white belly'
    # caption = 'this bird has a blue crown with white throat and brown secondaries'
    # caption = 'this bird has wings that are black and has a white belly'
    # caption = 'a yellow bird has wings with dark stripes and small eyes'
    # caption = 'a black bird has wings with dark stripes and small eyes'
    # caption = 'a red bird has wings with dark stripes and small eyes'
    # caption = 'a white bird has wings with dark stripes and small eyes'
    # caption = 'a blue bird has wings with dark stripes and small eyes'
    # caption = 'a pink bird has wings with dark stripes and small eyes'
    # caption = 'this is a white and grey bird with black wings and a black stripe by its eyes'
    # caption = 'a small bird with an orange bill and grey crown and breast'
    # caption = 'a small bird with black gray and white wingbars'
    # caption = 'this bird is white and light orange in color with a black beak'
    # caption = 'a small sized bird that has tones of brown and a short pointed bill' # beak?

    # MS coco captions
    # caption = 'two men skiing down a snow covered mountain in the evening'
    # caption = 'a man walking down a grass covered mountain'
    # caption = 'a close up of a boat on a field under a sunset'
    # caption = 'a close up of a boat on a field with a clear sky'
    # caption = 'a herd of black and white cattle standing on a field'
    # caption = 'a herd of black and white sheep standing on a field'
    # caption = 'a herd of black and white dogs standing on a field'
    # caption = 'a herd of brown cattle standing on a field'
    # caption = 'a herd of black and white cattle standing in a river'
    # caption = 'some horses in a field of green grass with a sky in the background'
    # caption = 'some horses in a field of yellow grass with a sky in the background'
    caption = 'some horses in a field of green grass with a sunset in the background'

    # convert caption to index
    caption_idx, caption_len = get_caption_idx(dataset_json, caption)
    caption_idx = torch.LongTensor(caption_idx)
    caption_len = torch.LongTensor([caption_len])
    caption_idx = caption_idx.view(1, -1)
    caption_len = caption_len.view(-1)

    # use rnn encoder to get caption embedding
    hidden = text_encoder.init_hidden(1)
    words_embs, sent_emb = text_encoder(caption_idx, caption_len, hidden)

    # generate fake image
    noise.data.normal_(0, 1)
    sent_emb = sent_emb.repeat(num_noise, 1)
    words_embs = words_embs.repeat(num_noise, 1, 1)
    with torch.no_grad():
        fake_imgs, fusion_mask = netG(noise, sent_emb)

        # create path to save image, caption and mask
        cap_number = 10000
        main_path = 'result/mani/cap_%s_0_coco_ch64' % (str(cap_number))
        img_save_path = '%s/image' % main_path
        mask_save_path = '%s/mask_' % main_path
        mkdir_p(img_save_path)
        for i in range(7):
            mkdir_p(mask_save_path + str(i))

        # save caption as image
        ixtoword = {v: k for k, v in dataset_json['word2idx'].items()}
        cap_img = cap2img(ixtoword, caption_idx, caption_len)
        im = cap_img[0].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)
        full_path = '%s/caption.png' % main_path
        im.save(full_path)

        # save generated images and masks
        for i in tqdm(range(num_noise)):
            full_path = '%s/image_%d.png' % (img_save_path, i)
            im = fake_imgs[i].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)
            im.save(full_path)

            for j in range(7):
                full_path = '%s%1d/mask_%d.png' % (mask_save_path, j, i)
                im = fusion_mask[j][i][0].data.cpu().numpy()
                im = im * 255
                im = im.astype(np.uint8)
                im = Image.fromarray(im)
                im.save(full_path)
Exemple #3
0
def train_network():

    init_epoch = 0
    best_f1 = 0
    total_steps = 0
    train_dir = ct.TRAIN_TXT
    val_dir = ct.VAL_TXT
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    torch.backends.cudnn.benchmark = True

    train_data = OSCD_TRAIN(train_dir)
    train_dataloader = DataLoader(train_data,
                                  batch_size=ct.BATCH_SIZE,
                                  shuffle=True)
    val_data = OSCD_TEST(val_dir)
    val_dataloader = DataLoader(val_data, batch_size=1, shuffle=False)
    netg = NetG(ct.ISIZE, ct.NC * 2, ct.NZ, ct.NDF,
                ct.EXTRALAYERS).to(device=device)
    netd = NetD(ct.ISIZE, ct.GT_C, 1, ct.NGF, ct.EXTRALAYERS).to(device=device)
    netg.apply(weights_init)
    netd.apply(weights_init)
    if ct.RESUME:
        assert os.path.exists(os.path.join(ct.WEIGHTS_SAVE_DIR, 'current_netG.pth')) \
                and os.path.exists(os.path.join(ct.WEIGHTS_SAVE_DIR, 'current_netG.pth')), \
                'There is not found any saved weights'
        print("\nLoading pre-trained networks.")
        init_epoch = torch.load(
            os.path.join(ct.WEIGHTS_SAVE_DIR, 'current_netG.pth'))['epoch']
        netg.load_state_dict(
            torch.load(os.path.join(ct.WEIGHTS_SAVE_DIR,
                                    'current_netG.pth'))['model_state_dict'])
        netd.load_state_dict(
            torch.load(os.path.join(ct.WEIGHTS_SAVE_DIR,
                                    'current_netD.pth'))['model_state_dict'])
        with open(os.path.join(ct.OUTPUTS_DIR, 'f1_score.txt')) as f:
            lines = f.readlines()
            best_f1 = float(lines[-2].strip().split(':')[-1])
        print("\tDone.\n")

    l_adv = l2_loss
    l_con = nn.L1Loss()
    l_enc = l2_loss
    l_bce = nn.BCELoss()
    l_cos = cos_loss
    dice = DiceLoss()
    optimizer_d = optim.Adam(netd.parameters(), lr=ct.LR, betas=(0.5, 0.999))
    optimizer_g = optim.Adam(netg.parameters(), lr=ct.LR, betas=(0.5, 0.999))

    start_time = time.time()
    for epoch in range(init_epoch + 1, ct.EPOCH):
        loss_g = []
        loss_d = []
        netg.train()
        netd.train()
        epoch_iter = 0
        for i, data in enumerate(train_dataloader):
            INPUT_SIZE = [ct.ISIZE, ct.ISIZE]
            x1, x2, gt = data
            x1 = x1.to(device, dtype=torch.float)
            x2 = x2.to(device, dtype=torch.float)
            gt = gt.to(device, dtype=torch.float)
            gt = gt[:, 0, :, :].unsqueeze(1)
            x = torch.cat((x1, x2), 1)

            epoch_iter += ct.BATCH_SIZE
            total_steps += ct.BATCH_SIZE
            real_label = torch.ones(size=(x1.shape[0], ),
                                    dtype=torch.float32,
                                    device=device)
            fake_label = torch.zeros(size=(x1.shape[0], ),
                                     dtype=torch.float32,
                                     device=device)

            #forward

            fake = netg(x)
            pred_real = netd(gt)
            pred_fake = netd(fake).detach()
            err_d_fake = l_bce(pred_fake, fake_label)
            err_g = l_con(fake, gt)
            err_g_total = ct.G_WEIGHT * err_g + ct.D_WEIGHT * err_d_fake

            pred_fake_ = netd(fake.detach())
            err_d_real = l_bce(pred_real, real_label)
            err_d_fake_ = l_bce(pred_fake_, fake_label)
            err_d_total = (err_d_real + err_d_fake_) * 0.5

            #backward
            optimizer_g.zero_grad()
            err_g_total.backward(retain_graph=True)
            optimizer_g.step()
            optimizer_d.zero_grad()
            err_d_total.backward()
            optimizer_d.step()

            errors = utils.get_errors(err_d_total, err_g_total)
            loss_g.append(err_g_total.item())
            loss_d.append(err_d_total.item())

            counter_ratio = float(epoch_iter) / len(train_dataloader.dataset)
            if (i % ct.DISPOLAY_STEP == 0 and i > 0):
                print(
                    'epoch:', epoch, 'iteration:', i,
                    ' G|D loss is {}|{}'.format(np.mean(loss_g[-51:]),
                                                np.mean(loss_d[-51:])))
                if ct.DISPLAY:
                    utils.plot_current_errors(epoch, counter_ratio, errors,
                                              vis)
                    utils.display_current_images(gt.data, fake.data, vis)
        utils.save_current_images(epoch, gt.data, fake.data, ct.IM_SAVE_DIR,
                                  'training_output_images')

        with open(os.path.join(ct.OUTPUTS_DIR, 'train_loss.txt'), 'a') as f:
            f.write(
                'after %s epoch, loss is %g,loss1 is %g,loss2 is %g,loss3 is %g'
                % (epoch, np.mean(loss_g), np.mean(loss_d), np.mean(loss_g),
                   np.mean(loss_d)))
            f.write('\n')
        if not os.path.exists(ct.WEIGHTS_SAVE_DIR):
            os.makedirs(ct.WEIGHTS_SAVE_DIR)
        utils.save_weights(epoch, netg, optimizer_g, ct.WEIGHTS_SAVE_DIR,
                           'netG')
        utils.save_weights(epoch, netd, optimizer_d, ct.WEIGHTS_SAVE_DIR,
                           'netD')
        duration = time.time() - start_time
        print('training duration is %g' % duration)

        #val phase
        print('Validating.................')
        pretrained_dict = torch.load(
            os.path.join(ct.WEIGHTS_SAVE_DIR,
                         'current_netG.pth'))['model_state_dict']
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        net = NetG(ct.ISIZE, ct.NC * 2, ct.NZ, ct.NDF,
                   ct.EXTRALAYERS).to(device=device)
        net.load_state_dict(pretrained_dict, False)
        with net.eval() and torch.no_grad():
            TP = 0
            FN = 0
            FP = 0
            TN = 0
            for k, data in enumerate(val_dataloader):
                x1, x2, label = data
                x1 = x1.to(device, dtype=torch.float)
                x2 = x2.to(device, dtype=torch.float)
                label = label.to(device, dtype=torch.float)
                label = label[:, 0, :, :].unsqueeze(1)
                x = torch.cat((x1, x2), 1)
                time_i = time.time()
                v_fake = net(x)

                tp, fp, tn, fn = eva.f1(v_fake, label)
                TP += tp
                FN += fn
                TN += tn
                FP += fp

            precision = TP / (TP + FP + 1e-8)
            oa = (TP + TN) / (TP + FN + TN + FP + 1e-8)
            recall = TP / (TP + FN + 1e-8)
            f1 = 2 * precision * recall / (precision + recall + 1e-8)
            if not os.path.exists(ct.BEST_WEIGHT_SAVE_DIR):
                os.makedirs(ct.BEST_WEIGHT_SAVE_DIR)
            if f1 > best_f1:
                best_f1 = f1
                shutil.copy(
                    os.path.join(ct.WEIGHTS_SAVE_DIR, 'current_netG.pth'),
                    os.path.join(ct.BEST_WEIGHT_SAVE_DIR, 'netG.pth'))
            print('current F1: {}'.format(f1))
            print('best f1: {}'.format(best_f1))
            with open(os.path.join(ct.OUTPUTS_DIR, 'f1_score.txt'), 'a') as f:
                f.write('current epoch:{},current f1:{},best f1:{}'.format(
                    epoch, f1, best_f1))
                f.write('\n')
class AnoGAN:
    """AnoGAN Class
    """
    def __init__(self, opt):
        # super(AnoGAN, self).__init__(opt, dataloader)

        # Initalize variables.
        self.opt = opt

        self.niter = self.opt.niter
        self.start_iter = 0
        self.netd_niter = 5
        self.test_iter = 100
        self.lr = self.opt.lr
        self.batchsize = {'train': self.opt.batchsize, 'test': 1}

        self.pretrained = False

        self.phase = 'train'
        self.outf = self.opt.experiment_group
        self.algorithm = 'wgan'

        # LOAD DATA SET
        self.dataloader = {
            'train':
            provider('train',
                     opt.category,
                     batch_size=self.batchsize['train'],
                     num_workers=4),
            'test':
            provider('test',
                     opt.category,
                     batch_size=self.batchsize['test'],
                     num_workers=4)
        }

        self.trn_dir = os.path.join(self.outf, self.opt.experiment_name,
                                    'train')
        self.tst_dir = os.path.join(self.outf, self.opt.experiment_name,
                                    'test')

        self.test_img_dir = os.path.join(self.outf, self.opt.experiment_name,
                                         'test', 'images')
        if not os.path.isdir(self.test_img_dir):
            os.makedirs(self.test_img_dir)

        self.best_test_dir = os.path.join(self.outf, self.opt.experiment_name,
                                          'test', 'best_images')
        if not os.path.isdir(self.best_test_dir):
            os.makedirs(self.best_test_dir)

        self.weight_dir = os.path.join(self.trn_dir, 'weights')
        if not os.path.exists(self.weight_dir): os.makedirs(self.weight_dir)

        # -- Misc attributes
        self.epoch = 0

        self.l_con = l1_loss
        self.l_enc = l2_loss

        ##
        # Create and initialize networks.
        self.netg = NetG().cuda()
        self.netd = NetD().cuda()

        # Setup optimizer
        self.optimizer_d = optim.RMSprop(self.netd.parameters(), lr=self.lr)
        self.optimizer_g = optim.Adam(self.netg.parameters(), lr=self.lr)

        ##
        self.weight_path = os.path.join(self.outf, self.opt.experiment_name,
                                        'train', 'weights')
        if os.path.exists(self.weight_path) and len(
                os.listdir(self.weight_path)) == 2:
            print("Loading pre-trained networks...\n")
            self.netg.load_state_dict(
                torch.load(os.path.join(self.weight_path,
                                        'netG.pth'))['state_dict'])
            self.netd.load_state_dict(
                torch.load(os.path.join(self.weight_path,
                                        'netD.pth'))['state_dict'])

            self.optimizer_g.load_state_dict(
                torch.load(os.path.join(self.weight_path,
                                        'netG.pth'))['optimizer'])
            self.optimizer_d.load_state_dict(
                torch.load(os.path.join(self.weight_path,
                                        'netD.pth'))['optimizer'])

            self.start_iter = torch.load(
                os.path.join(self.weight_path, 'netG.pth'))['epoch']

    ##
    def start(self):
        """ Train the model
        """

        ##
        # TRAIN
        # self.total_steps = 0
        best_criterion = -1  #float('inf')
        best_auc = -1

        # Train for niter epochs.
        # print(">> Training model %s." % self.name)
        for self.epoch in range(self.start_iter, self.niter):
            # Train for one epoch
            mean_wass = self.train()

            (auc, res, best_rec, best_threshold), res_total = self.test()
            message = ''
            # message += 'criterion: (%.3f+%.3f)/2=%.3f ' % (best_rec[0], best_rec[1], res)
            # message += 'best threshold: %.3f ' % best_threshold
            message += 'Wasserstein Distance:%.3d ' % mean_wass
            message += 'AUC: %.3f ' % auc

            print(message)

            torch.save(
                {
                    'epoch': self.epoch + 1,
                    'state_dict': self.netg.state_dict(),
                    'optimizer': self.optimizer_g.state_dict()
                }, '%s/netG.pth' % (self.weight_dir))

            torch.save(
                {
                    'epoch': self.epoch + 1,
                    'state_dict': self.netd.state_dict(),
                    'optimizer': self.optimizer_d.state_dict()
                }, '%s/netD.pth' % (self.weight_dir))

            if auc > best_auc:
                best_auc = auc
                new_message = "******** New optimal found, saving state ********"
                message = message + '\n' + new_message
                print(new_message)

                for img in os.listdir(self.best_test_dir):
                    os.remove(os.path.join(self.best_test_dir, img))

                for img in os.listdir(self.test_img_dir):
                    shutil.copyfile(os.path.join(self.test_img_dir, img),
                                    os.path.join(self.best_test_dir, img))

                shutil.copyfile('%s/netG.pth' % (self.weight_dir),
                                '%s/netg_best.pth' % (self.weight_dir))

            log_name = os.path.join(self.outf, self.opt.experiment_name,
                                    'loss_log.txt')
            message = 'Epoch%3d:' % self.epoch + ' ' + message
            with open(log_name, "a") as log_file:
                if self.epoch == 0:
                    log_file.write('\n\n')
                log_file.write('%s\n' % message)

        print(">> Training %s Done..." % self.opt.experiment_name)

    ##
    def train(self):
        """ Train the model for one epoch.
        """
        print("\n>>> Epoch %d/%d, Running " % (self.epoch + 1, self.niter) +
              self.opt.experiment_name)

        self.netg.train()
        self.netd.train()
        # for p in self.netg.parameters(): p.requires_grad = True

        mean_wass = 0

        tk0 = tqdm(self.dataloader['train'],
                   total=len(self.dataloader['train']))
        for i, itr in enumerate(tk0):
            input, _ = itr
            input = input.cuda()
            wasserstein_d = None
            # if self.algorithm == 'wgan':
            # train NetD
            for _ in range(self.netd_niter):
                # for p in self.netd.parameters(): p.requires_grad = True
                self.optimizer_d.zero_grad()

                # forward_g
                latent_i = torch.rand(self.batchsize['train'], 64, 1, 1).cuda()
                fake = self.netg(latent_i)
                # forward_d
                _, pred_real = self.netd(input)
                _, pred_fake = self.netd(fake)  # .detach() TODO

                # Backward-pass
                wasserstein_d = (pred_fake.mean() - pred_real.mean()) * 1
                wasserstein_d.backward()
                self.optimizer_d.step()

                for p in self.netd.parameters():
                    p.data.clamp_(-0.01, 0.01)  #<<<<<<<

            # train netg
            # for p in self.netd.parameters(): p.requires_grad = False
            self.optimizer_g.zero_grad()
            noise = torch.rand(self.batchsize['train'], 64, 1, 1).cuda()
            fake = self.netg(noise)
            _, pred_fake = self.netd(fake)
            err_g_d = -pred_fake.mean()  # negative

            err_g_d.backward()
            self.optimizer_g.step()

            errors = {
                'loss_netD': wasserstein_d.item(),
                'loss_netG': round(err_g_d.item(), 3),
            }

            mean_wass += wasserstein_d.item()
            tk0.set_postfix(errors)

            if i % 50 == 0:
                img_dir = os.path.join(self.outf, self.opt.experiment_name,
                                       'train', 'images')
                if not os.path.isdir(img_dir):
                    os.makedirs(img_dir)
                self.save_image_cv2(input.data, '%s/reals.png' % img_dir)
                self.save_image_cv2(fake.data,
                                    '%s/fakes%03d.png' % (img_dir, i))

        mean_wass /= len(self.dataloader['train'])
        return mean_wass

    ##
    def test(self):
        """ Test AnoGAN model.

        Args:
            dataloader ([type]): Dataloader for the test set

        Raises:
            IOError: Model weights not found.
        """
        self.netg.eval()
        self.netd.eval()
        # for p in self.netg.parameters(): p.requires_grad = False
        # for p in self.netd.parameters(): p.requires_grad = False

        for img in os.listdir(self.test_img_dir):
            os.remove(os.path.join(self.test_img_dir, img))

        self.phase = 'test'
        meter = Meter_AnoGAN()
        tk1 = tqdm(self.dataloader['test'], total=len(self.dataloader['test']))
        for i, itr in enumerate(tk1):
            input, target = itr
            input = input.cuda()

            latent_i = torch.rand(self.batchsize['test'], 64, 1, 1).cuda()
            latent_i.requires_grad = True

            optimizer_latent = optim.Adam([latent_i], lr=self.lr)
            test_loss = None
            for _ in range(self.test_iter):
                optimizer_latent.zero_grad()
                fake = self.netg(latent_i)
                residual_loss = self.l_con(input, fake)
                latent_o, _ = self.netd(fake)
                discrimination_loss = self.l_enc(latent_i, latent_o)
                alpha = 0.1
                test_loss = (
                    1 - alpha) * residual_loss + alpha * discrimination_loss
                test_loss.backward()
                optimizer_latent.step()

            abnormal_score = test_loss
            meter.update(abnormal_score, target)  #<<<TODO

            # Save test images.
            combine = torch.cat([input.cpu(), fake.cpu()], dim=0)
            self.save_image_cv2(combine,
                                '%s/%05d.jpg' % (self.test_img_dir, i + 1))

        criterion, res_total = meter.get_metrics()

        # rename images
        for i, res in enumerate(res_total):
            os.rename('%s/%05d.jpg' % (self.test_img_dir, i + 1),
                      '%s/%05d_%s.jpg' % (self.test_img_dir, i + 1, res))

        return criterion, res_total

    @staticmethod
    def save_image_cv2(tensor, filename):
        # return
        from torchvision.utils import make_grid
        # tensor = (tensor + 1) / 2
        grid = make_grid(tensor, 8, 2, 0, False, None, False)
        ndarray = grid.mul_(255).clamp_(0, 255).permute(1, 2, 0).to(
            'cpu', torch.uint8).numpy()
        cv2.imwrite(filename, ndarray)
Exemple #5
0
def test_network():
    threshold = ct.THRESHOLD
    test_dir = ct.TEST_TXT
    path = os.path.join(ct.BEST_WEIGHT_SAVE_DIR, 'netG.pth')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    pretrained_dict = torch.load(
        path, map_location=torch.device(device))['model_state_dict']
    test_data = OSCD_TEST(test_dir)
    test_dataloader = DataLoader(test_data, batch_size=1, shuffle=False)
    net = NetG(ct.ISIZE, ct.NC * 2, ct.NZ, ct.NDF, ct.EXTRALAYERS).to(device)
    #     net = nn.DataParallel(net)
    net.load_state_dict(pretrained_dict, False)
    torch.no_grad()
    net.eval()
    i = 0
    TP = 0
    FN = 0
    FP = 0
    TN = 0
    for i, data in enumerate(test_dataloader):
        INPUT_SIZE = [ct.ISIZE, ct.ISIZE]
        x1, x2, gt = data
        x1 = x1.to(device, dtype=torch.float)
        x2 = x2.to(device, dtype=torch.float)
        gt = gt.to(device, dtype=torch.float)
        gt = gt[:, 0, :, :].unsqueeze(1)

        x = torch.cat((x1, x2), 1)
        fake = net(x)

        save_path = os.path.join(ct.IM_SAVE_DIR, 'test_output_images')
        if not os.path.isdir(save_path):
            os.makedirs(save_path)

        if ct.SAVE_TEST_IAMGES:
            vutils.save_image(x1.data,
                              os.path.join(save_path, '%d_x1.png' % i),
                              normalize=True)
            vutils.save_image(x2.data,
                              os.path.join(save_path, '%d_x2.png' % i),
                              normalize=True)
            vutils.save_image(fake.data,
                              os.path.join(save_path, '%d_gt_fake.png' % i),
                              normalize=True)
            vutils.save_image(gt,
                              os.path.join(save_path, '%d_gt.png' % i),
                              normalize=True)

        tp, fp, tn, fn = eva.f1(fake, gt)
        TP += tp
        FN += fn
        TN += tn
        FP += fp
        i += 1
        print('testing {}th images'.format(i))
    iou = TP / (FN + TP + FP + 1e-8)
    precision = TP / (TP + FP + 1e-8)
    oa = (TP + TN) / (TP + FN + TN + FP + 1e-8)
    recall = TP / (TP + FN + 1e-8)
    f1 = 2 * precision * recall / (precision + recall + 1e-8)
    P = ((TP + FP) * (TP + FN) + (FN + TN) *
         (FP + TN)) / ((TP + TN + FP + FN)**2 + 1e-8)
    Kappa = (oa - P) / (1 - P + 1e-8)
    results = {
        'iou': iou,
        'precision': precision,
        'oa': oa,
        'recall': recall,
        'f1': f1,
        'kappa': Kappa
    }

    with open(os.path.join(ct.OUTPUTS_DIR, 'test_score.txt'), 'a') as f:
        f.write('-----test results on the best model {}-----'.format(
            time.strftime('%Y-%m-%d %H:%M:%S')))
        f.write('\n')
        for key, value in results.items():
            print(key, value)
            f.write('{}: {}'.format(key, value))
            f.write('\n')