YDTR/Test.py at main · tthinking/YDTR · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
from PIL import Image
import numpy as np
import os
import torch

import time
import imageio

import torchvision.transforms as transforms

from Networks.net import MODEL as net

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

device = torch.device('cuda:0')


model = net(in_channel=1)


model_path = "./models/model_10.pth"
use_gpu = torch.cuda.is_available()

if use_gpu:

    model = model.cuda()
    model.cuda()
    model.load_state_dict(torch.load(model_path))

else:

    state_dict = torch.load(model_path, map_location='cpu')
    model.load_state_dict(state_dict)


def fusion():
    tic = time.time()
    for num in range(1):

        path1 = './test_images/IR_002.bmp'
        path2 = './test_images/VIS_002.bmp'
        img1 = Image.open(path1).convert('L')
        img2 = Image.open(path2).convert('L')

        img1_org = img1
        img2_org = img2
        tran = transforms.ToTensor()
        img1_org = tran(img1_org)
        img2_org = tran(img2_org)
        if use_gpu:
            img1_org  = img1_org.cuda()
            img2_org = img2_org.cuda()
        else:
            img1_org = img1_org
            img2_org = img2_org
        img1_org = img1_org.unsqueeze(0)
        img2_org = img2_org.unsqueeze(0)

        model.eval()
        out = model(img1_org, img2_org )

        d = np.squeeze(out.detach().cpu().numpy())

        result = (d* 255).astype(np.uint8)


        imageio.imwrite('./fusion result/{}.bmp'.format( num), result)

    toc = time.time()
    print('end {}{}'.format(num // 10, num % 10), ', time:{}'.format(toc - tic))


if __name__ == '__main__':

    fusion()