File size: 807 Bytes
7d1312d |
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 |
import numpy as np
from PIL import Image
from torchvision import transforms
def requires_grad(model, flag=True):
for p in model.parameters():
p.requires_grad = flag
def get_keys(d, name):
if 'state_dict' in d:
d = d['state_dict']
d_filt = {k[len(name) + 1:]: v for k, v in d.items() if k[:len(name)] == name}
return d_filt
def load_img(path_img, img_size=(256, 256)):
transform = transforms.Compose(
[transforms.Resize(img_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))])
if type(path_img) is np.ndarray:
img = Image.fromarray(path_img)
else:
img = Image.open(path_img).convert('RGB')
img = transform(img)
img.unsqueeze_(0)
return img
|