自定义的图像操作库 mii

自定义的图像操作库 mii，放在 环境\Lib\site-packages 中
如Anaconda3\envs\pytorch\Lib\site-packages\mii

mii
|——__init__.py
|——tools.py

__init__.py ：

from mii.tools import * 

tools.py:

import cv2
import torch
import numpy as np
from einops import rearrange


def imread(img_path, tensor=False, shape="hwc"):
    img = cv2.imread(img_path)
    img = bgr2rgb(img)
    img = np.array(img, dtype=np.float32)
    if tensor == True:
        img = ndarray2tensor(img)
    if shape == "chw":
        img = hwc2chw(img)
    return img


def imwrite(img, file_name):
    if type(img) == torch.Tensor:
        img = tensor2ndarray(img)
    img = np.array(img, dtype=np.uint8)
    img = bgr2rgb(img)
    cv2.imwrite(file_name, img)


def hwc2chw(img):
    if len(img.shape) == 3:
        img = rearrange(img, "h w c -> c h w")
    elif len(img.shape) == 4:
        img = rearrange(img, "b h w c -> b c h w")
    return img


def chw2hwc(img):
    if len(img.shape) == 3:
        img = rearrange(img, "c h w -> h w c")
    elif len(img.shape) == 4:
        img = rearrange(img, "b c h w -> b h w c")
    return img


def bgr2rgb(img):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img


def ndarray2tensor(img, dytpe=torch.float32):
    img = torch.tensor(img, dtype=dytpe)
    return img


def tensor2ndarray(img, dtype=np.float32):
    img = img.cpu().detach().numpy()
    img = np.array(img, dtype=dtype)
    return img


def read_dir(dir, s=1, e=256,tensor=False, shape="hwc"):
    img_s = imread(f"{dir}/{s}.png")
    h,w,c = img_s.shape

    img_batch = np.zeros([e-s+1,h,w,c])
    for i in range(s, e + 1):
        img = imread(f"{dir}/{i}.png")
        if img.shape != img_s.shape:
            img = cv2.resize(img,[w,h],interpolation=cv2.INTER_CUBIC)
        img_batch[i-s] = img
    
    if tensor == True:
        img_batch = ndarray2tensor(img_batch)
    
    if shape == "chw":
        img_batch = hwc2chw(img_batch)
    
    return img_batch