210 lines
7.7 KiB
Python
210 lines
7.7 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Part of the implementation is borrowed and modified from ControlNet,
|
|
# publicly available at https://github.com/lllyasviel/ControlNet
|
|
|
|
import math
|
|
import os
|
|
from typing import Any, Dict
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
|
|
from modelscope.metainfo import Preprocessors
|
|
from modelscope.models.cv.controllable_image_generation.annotator.annotator import (
|
|
CannyDetector, HEDdetector, MidasDetector, MLSDdetector, OpenposeDetector,
|
|
SegformerDetector, nms)
|
|
from modelscope.preprocessors.base import Preprocessor
|
|
from modelscope.preprocessors.builder import PREPROCESSORS
|
|
from modelscope.preprocessors.image import load_image
|
|
from modelscope.utils.constant import (DEFAULT_MODEL_REVISION, Fields, Invoke,
|
|
ModeKeys, Tasks)
|
|
from modelscope.utils.type_assert import type_assert
|
|
|
|
|
|
def HWC3(x):
|
|
assert x.dtype == np.uint8
|
|
if x.ndim == 2:
|
|
x = x[:, :, None]
|
|
assert x.ndim == 3
|
|
H, W, C = x.shape
|
|
assert C == 1 or C == 3 or C == 4
|
|
if C == 3:
|
|
return x
|
|
if C == 1:
|
|
return np.concatenate([x, x, x], axis=2)
|
|
if C == 4:
|
|
color = x[:, :, 0:3].astype(np.float32)
|
|
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
|
|
y = color * alpha + 255.0 * (1.0 - alpha)
|
|
y = y.clip(0, 255).astype(np.uint8)
|
|
return y
|
|
|
|
|
|
def resize_image(input_image, resolution):
|
|
H, W, C = input_image.shape
|
|
H = float(H)
|
|
W = float(W)
|
|
k = float(resolution) / min(H, W)
|
|
H *= k
|
|
W *= k
|
|
H = int(np.round(H / 64.0)) * 64
|
|
W = int(np.round(W / 64.0)) * 64
|
|
img = cv2.resize(
|
|
input_image, (W, H),
|
|
interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
|
|
return img
|
|
|
|
|
|
def build_detector(control_type, model_path, device):
|
|
if control_type == 'scribble':
|
|
detector = None
|
|
elif control_type == 'canny':
|
|
detector = CannyDetector()
|
|
elif control_type == 'hough':
|
|
detector = MLSDdetector(model_path, device)
|
|
elif control_type == 'hed':
|
|
detector = HEDdetector(model_path, device)
|
|
elif control_type == 'depth':
|
|
detector = MidasDetector(model_path, device)
|
|
elif control_type == 'normal':
|
|
detector = MidasDetector(model_path, device)
|
|
elif control_type == 'pose':
|
|
detector = OpenposeDetector(model_path, device)
|
|
elif control_type == 'seg':
|
|
detector = SegformerDetector(model_path, device)
|
|
elif control_type == 'fake_scribble':
|
|
detector = HEDdetector(model_path, device)
|
|
else:
|
|
detector = HEDdetector(model_path, device)
|
|
return detector
|
|
|
|
|
|
def get_detected_map(detector, control_type, img, **kwargs):
|
|
if control_type == 'scribble':
|
|
detected_map = np.zeros_like(img, dtype=np.uint8)
|
|
detected_map[np.min(img, axis=2) < 127] = 255
|
|
elif control_type == 'canny':
|
|
detected_map = detector(img, kwargs['low_threshold'],
|
|
kwargs['high_threshold'])
|
|
detected_map = HWC3(detected_map)
|
|
elif control_type == 'hough':
|
|
detected_map = detector(img, kwargs['value_threshold'],
|
|
kwargs['distance_threshold'])
|
|
detected_map = HWC3(detected_map)
|
|
elif control_type == 'hed':
|
|
detected_map = detector(img)
|
|
detected_map = HWC3(detected_map)
|
|
elif control_type == 'depth':
|
|
H, W, C = img.shape
|
|
det_img = resize_image(img, 384)
|
|
detected_map, _ = detector(det_img)
|
|
detected_map = HWC3(detected_map)
|
|
detected_map = cv2.resize(
|
|
detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
|
|
elif control_type == 'normal':
|
|
H, W, C = img.shape
|
|
det_img = resize_image(img, 384)
|
|
_, detected_map = detector(det_img, bg_th=kwargs['bg_threshold'])
|
|
detected_map = HWC3(detected_map)
|
|
detected_map = cv2.resize(
|
|
detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
|
|
detected_map = detected_map[:, :, ::-1]
|
|
elif control_type == 'pose':
|
|
detected_map, _ = detector(img)
|
|
detected_map = HWC3(detected_map)
|
|
elif control_type == 'seg':
|
|
detected_map = detector(img)
|
|
elif control_type == 'fake_scribble':
|
|
detected_map = detector(img)
|
|
detected_map = HWC3(detected_map)
|
|
detected_map = nms(detected_map, 127, 3.0)
|
|
detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
|
|
detected_map[detected_map > 4] = 255
|
|
detected_map[detected_map < 255] = 0
|
|
|
|
return detected_map
|
|
|
|
|
|
@PREPROCESSORS.register_module(
|
|
Fields.cv,
|
|
module_name=Preprocessors.controllable_image_generation_preprocessor)
|
|
class ControllableImageGenerationPreprocessor(Preprocessor):
|
|
|
|
def __init__(self, mode=ModeKeys.INFERENCE, *args, **kwargs):
|
|
super().__init__(mode=ModeKeys.INFERENCE, *args, **kwargs)
|
|
self.detector = build_detector(
|
|
kwargs.get('control_type', 'hed'), kwargs.get('model_path', None),
|
|
kwargs.get('device', 'cuda'))
|
|
|
|
@type_assert(object, object)
|
|
def __call__(self, data: input, **kwargs) -> Dict[str, Any]:
|
|
image_resolution = data.get('image_resolution',
|
|
kwargs['image_resolution'])
|
|
image = np.array(load_image(data['image']))
|
|
image = resize_image(HWC3(image), image_resolution)
|
|
print(f'Test with image resolution: {image_resolution}')
|
|
|
|
is_cat_img = kwargs.get('is_cat_img', True)
|
|
|
|
if 'prompt' in data.keys():
|
|
model_prompt = data['prompt']
|
|
else:
|
|
# for demo_service
|
|
model_prompt = kwargs.get('prompt', '')
|
|
print(f'Test with prompt: {model_prompt}')
|
|
|
|
control_type = kwargs.get('control_type', 'hed')
|
|
print(f'Test with input type: {control_type}')
|
|
|
|
save_memory = kwargs.get('save_memory', False)
|
|
|
|
# generate detected_map
|
|
if control_type == 'scribble':
|
|
detected_map = get_detected_map(self.detector, control_type, image)
|
|
elif control_type == 'canny':
|
|
low_threshold = kwargs['modelsetting'].canny.low_threshold
|
|
high_threshold = kwargs['modelsetting'].canny.high_threshold
|
|
detected_map = get_detected_map(
|
|
self.detector,
|
|
control_type,
|
|
image,
|
|
low_threshold=low_threshold,
|
|
high_threshold=high_threshold)
|
|
elif control_type == 'hough':
|
|
value_threshold = kwargs['modelsetting'].hough.value_threshold
|
|
distance_threshold = kwargs[
|
|
'modelsetting'].hough.distance_threshold
|
|
detected_map = get_detected_map(
|
|
self.detector,
|
|
control_type,
|
|
image,
|
|
value_threshold=value_threshold,
|
|
distance_threshold=distance_threshold)
|
|
elif control_type in ['hed', 'depth', 'pose', 'seg', 'fake_scribble']:
|
|
detected_map = get_detected_map(self.detector, control_type, image)
|
|
elif control_type == 'normal':
|
|
bg_threshold = kwargs['modelsetting'].normal.bg_threshold
|
|
detected_map = get_detected_map(
|
|
self.detector, control_type, image, bg_threshold=bg_threshold)
|
|
else:
|
|
detected_map = get_detected_map(
|
|
self.detector, control_type='hed', img=image)
|
|
|
|
input_dict = {
|
|
'image': image,
|
|
'prompt': model_prompt,
|
|
'detected_map': detected_map,
|
|
'save_memory': save_memory,
|
|
'is_cat_img': is_cat_img
|
|
}
|
|
|
|
for k in data.keys():
|
|
if k not in input_dict.keys():
|
|
input_dict[k] = data[k]
|
|
|
|
return input_dict
|