Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Different demo result #16

Open
rudylyh opened this issue Jan 16, 2020 · 0 comments
Open

Different demo result #16

rudylyh opened this issue Jan 16, 2020 · 0 comments

Comments

@rudylyh
Copy link

rudylyh commented Jan 16, 2020

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant