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