Wan2GP/preprocessing/gray.py

# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.

import cv2
import numpy as np
from PIL import Image
import torch

def convert_to_numpy(image):
    if isinstance(image, Image.Image):
        image = np.array(image)
    elif isinstance(image, torch.Tensor):
        image = image.detach().cpu().numpy()
    elif isinstance(image, np.ndarray):
        image = image.copy()
    else:
        raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
    return image

class GrayAnnotator:
    def __init__(self, cfg):
        pass
    def forward(self, image):
        image = convert_to_numpy(image)
        gray_map = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        return gray_map[..., None].repeat(3, axis=2)


class GrayVideoAnnotator(GrayAnnotator):
    def forward(self, frames):
        ret_frames = []
        for frame in frames:
            anno_frame = super().forward(np.array(frame))
            ret_frames.append(anno_frame)
        return ret_frames