-
Notifications
You must be signed in to change notification settings - Fork 4
/
csp_test_pytorch.py
91 lines (75 loc) · 3.12 KB
/
csp_test_pytorch.py
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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import torch
import torchvision
from torchvision import transforms
import cv2
import numpy as np
import matplotlib.pyplot as plt
from steerable.SCFpyr_PyTorch import SCFpyr_PyTorch
import steerable.utils as utils
from net.phasenet import Triplets,show_Triplets_batch
# # Load dataset
# batch_size = 4
# transform = transforms.Compose(
# [transforms.Resize((256, 256)),
# transforms.ToTensor()])
# dataset = Triplets(
# '/home/lj/Documents/code/python/DAVIS/JPEGImages/480p/', transform)
# trainloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
# shuffle=True, num_workers=4)
# dataiter = iter(trainloader)
# Triplets_batch = dataiter.next()
# # get images_list [[N,C,H,W],[N,C,H,W],...], usually len(images_list)=batch_size if len(dataset)%batch_size==0
# images_list = [torch.stack([Triplets_batch['start'][i],
# Triplets_batch['inter'][i],
# Triplets_batch['end'][i]]) for i in range(len(Triplets_batch['start']))]
# # show_Triplets_batch(Triplets_batch)
# # Requires PyTorch with MKL when setting to 'cpu'
# device = torch.device('cpu')#cuda:0
# # Load batch of images [N,1,H,W]
# # im_batch_numpy = utils.load_image_batch('./assets/lena.jpg',32,200)
# # im_batch_torch = torch.from_numpy(im_batch_numpy).to(device)
# # Initialize Complex Steerbale Pyramid
# height = 12
# nbands = 4
# scale_factor = 2**(1/2)
# pyr = SCFpyr_PyTorch(height=height, nbands=nbands, scale_factor=scale_factor, device=device)
# pyr_type = 1
# # Decompose entire batch of images
# coeff = pyr.build(images_list[0][:,0,:,:].unsqueeze(1).to(device), pyr_type=pyr_type)
# import pdb;pdb.set_trace()
# # Reconstruct batch of images again
# batch_recon = pyr.reconstruct(coeff,pyr_type=pyr_type)
# cv2.imshow('img_recon',(batch_recon[0].cpu().numpy()*255).astype(np.uint8))
# # # Visualization
# # coeff_single = utils.extract_from_batch(coeff, 0)
# # coeff_grid = utils.make_grid_coeff(coeff_single, normalize=True)
# # cv2.imshow('Complex Steerable Pyramid', coeff_grid)
# cv2.waitKey(0)
# Requires PyTorch with MKL when setting to 'cpu'
device = torch.device('cpu')
# Load batch of images [N,1,H,W]
im_batch_numpy = utils.load_image_batch('./assets/lena.jpg',32,600)
img=cv2.imread('./assets/lena.jpg',0)
cv2.imshow('yuantu',img)
im_torch = torch.from_numpy(img).to(device)
im_batch_torch=im_torch.unsqueeze(0).unsqueeze(0).float()
# Initialize Complex Steerbale Pyramid
height = 12
nbands = 4
scale_factor = 2**(1/2)
pyr = SCFpyr_PyTorch(height=height, nbands=nbands, scale_factor=scale_factor, device=device)
pyr_type = 1
# Decompose entire batch of images
coeff = pyr.build(im_batch_torch,pyr_type)
# Reconstruct batch of images again
img_recon = pyr.reconstruct(coeff,pyr_type)
img=im_torch.float()
recon=img_recon.squeeze()
loss=torch.nn.MSELoss()
print('MSE:',loss(img,recon))
cv2.imshow('recon',recon.numpy().astype(np.uint8))
# Visualization
# coeff_single = utils.extract_from_batch(coeff, 0)
# coeff_grid = utils.make_grid_coeff(coeff_single, normalize=True)
# cv2.imshow('Complex Steerable Pyramid', coeff_grid)
cv2.waitKey(0)