utils.py 765 Bytes
'''
 _*_coding:utf-8 _*_
 @Time     :2022/1/28   19:58
 @Author   : qiaofengsheng
 @File     :utils.py
 @Software :PyCharm
 '''
import torch
import yaml
from PIL import Image
from torch.nn.functional import one_hot


def load_config_util(config_path):
    config_file = open(config_path, 'r', encoding='utf-8')
    config_data = yaml.load(config_file)
    return config_data


def keep_shape_resize(frame, size=128):
    w, h = frame.size
    temp = max(w, h)
    mask = Image.new('RGB', (temp, temp), (0, 0, 0))
    if w >= h:
        position = (0, (w - h) // 2)
    else:
        position = ((h - w) // 2, 0)
    mask.paste(frame, position)
    mask = mask.resize((size, size))
    return mask


def label_one_hot(label):
    return one_hot(torch.tensor(label))