ignacio-rocco / cnngeometric_pytorch

CNNGeometric PyTorch implementation

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Different demo result

rudylyh opened this issue · comments

I ran the demo script to test the provided "Affine+VGG+StreetView" model on the demo duck image pair, but the result affine image is not as good as it is in the ReadMe. Could anyone tell me what is the problem? Thanks a lot!

I simplified the demo script to the following:

from __future__ import print_function, division
import os
import argparse
from torch.utils.data import DataLoader
from model.cnn_geometric_model import CNNGeometric
from data.pf_dataset import PFDataset
from image.normalization import NormalizeImageDict, normalize_image
from util.torch_util import BatchTensorToVars
from geotnf.transformation import GeometricTnf
from collections import OrderedDict
import cv2

use_cuda = True
model_aff = CNNGeometric(use_cuda=use_cuda, feature_extraction_cnn='vgg')
checkpoint = torch.load('trained_models/best_streetview_checkpoint_adam_affine_grid_loss.pth.tar',
                        map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()])
model_aff.load_state_dict(checkpoint['state_dict'])

dataset = PFDataset(csv_file='datasets/PF-dataset/test_pairs_pf.csv', dataset_path='datasets',
                    transform=NormalizeImageDict(['source_image', 'target_image']))
dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4)
batchTensorToVars = BatchTensorToVars(use_cuda=use_cuda)
affTnf = GeometricTnf(geometric_model='affine', use_cuda=use_cuda)

for i, batch in enumerate(dataloader):
    batch = batchTensorToVars(batch)
    model_aff.eval()
    theta_aff = model_aff(batch)
    warped_image_aff = affTnf(batch['source_image'], theta_aff.view(-1, 2, 3))

    source_image = normalize_image(batch['source_image'], forward=False)
    source_image = source_image.data.squeeze(0).transpose(0, 1).transpose(1, 2).cpu().numpy()
    target_image = normalize_image(batch['target_image'], forward=False)
    target_image = target_image.data.squeeze(0).transpose(0, 1).transpose(1, 2).cpu().numpy()
    warped_image_aff = normalize_image(warped_image_aff, forward=False)
    warped_image_aff = warped_image_aff.data.squeeze(0).transpose(0, 1).transpose(1, 2).cpu().numpy()

    source_image = cv2.cvtColor(source_image, cv2.COLOR_RGB2BGR)
    target_image = cv2.cvtColor(target_image, cv2.COLOR_RGB2BGR)
    warped_image_aff = cv2.cvtColor(warped_image_aff, cv2.COLOR_RGB2BGR)

    cv2.imwrite('res_img/'+str(i)+'_source.png', (255*source_image).astype(np.uint8))
    cv2.imwrite('res_img/'+str(i)+'_target.png', (255*target_image).astype(np.uint8))
    cv2.imwrite('res_img/'+str(i)+'_result_aff.png', (255*warped_image_aff).astype(np.uint8))

    res = input('Run for another example ([y]/n): ')
    if res == 'n':
        break

The source, target and affined image respectively are:
0_source
0_target
0_result_aff