diff --git a/anyline.py b/anyline.py index 8c7a74c..b9af680 100644 --- a/anyline.py +++ b/anyline.py @@ -5,7 +5,10 @@ import kornia as kn from pathlib import Path from .teed_module import MTEED from huggingface_hub import hf_hub_download -import cv2 # Ensure you have OpenCV installed +import cv2 +from skimage import morphology +from PIL import Image +from einops import rearrange class AnyLine: @classmethod @@ -25,44 +28,189 @@ class AnyLine: def __init__(self): self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.model = self.load_model() + self.mteed_detector = self.MTEEDDetector(self.model, self.device) def load_model(self): - checkpoint_filename = "MTEED.safetetensors" + checkpoint_filename = "MTEED.pth" here = Path(__file__).parent.resolve() checkpoint_dir = here / "checkpoints" checkpoint_path = checkpoint_dir / checkpoint_filename + # Download the model if it's not present if not checkpoint_path.is_file(): print("Model not found locally, downloading from HuggingFace...") checkpoint_dir.mkdir(parents=True, exist_ok=True) - checkpoint_path = hf_hub_download(repo_id="TheMistoAI/MistoLine", filename=checkpoint_filename, cache_dir=checkpoint_dir) + checkpoint_path = hf_hub_download(repo_id="TheMistoAI/MistoLine", filename=checkpoint_filename, local_dir=checkpoint_dir) model = MTEED().to(self.device) model.load_state_dict(torch.load(checkpoint_path, map_location=self.device)) model.eval() return model - - def download_model(self): - hf_hub_download(repo_id=self.REPO_ID, filename=self.FILENAME, local_dir="./checkpoints", force_download=True) - def get_anyline(self, image): - image_tensor = torch.tensor(image).to(self.device) - original_tensor = image_tensor.clone() + # Process the image with MTEED model + mteed_result = self.mteed_detector(image[0], detect_resolution=image[0].shape[0]) + print("mteed_result: ", mteed_result) - with torch.no_grad(): - preds = self.model(image_tensor, original=original_tensor, is_eval=True) - img_height, img_width = image.shape[2], image.shape[3] - image_vis = kn.utils.tensor_to_image(torch.sigmoid(preds[0])) - image_vis = (255.0 * (1.0 - image_vis)).astype(np.uint8) - image_vis = cv2.resize(image_vis, (img_width, img_height)) + # Process the image with the lineart standard preprocessor + lineart_result = process_line_art(image[0]) + print("lineart_result: ", lineart_result) + + # Combine the results - return (image_vis,) -NODE_CLASS_MAPPINGS = { - "AnyLine": AnyLine -} + return (torch.tensor(lineart_result),) + + class MTEEDDetector: + def __init__(self, model, device): + self.model = model + self.device = device + def __call__(self, input_image, detect_resolution=512, safe_steps=2, upscale_method="INTER_CUBIC", output_type=None, **kwargs): + input_image, output_type = common_input_validate(input_image, output_type, **kwargs) + input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method) -NODE_DISPLAY_NAME_MAPPINGS = { - "AnyLine": "TheMisto Anyline" -} + H, W, _ = input_image.shape + with torch.no_grad(): + image_teed = torch.from_numpy(input_image.copy()).float().to(self.device) + image_teed = rearrange(image_teed, 'h w c -> 1 c h w') + edges = self.model(image_teed) + edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges] + edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges] + edges = np.stack(edges, axis=2) + edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64))) + if safe_steps != 0: + edge = safe_step(edge, safe_steps) + edge = (edge * 255.0).clip(0, 255).astype(np.uint8) + + detected_map = remove_pad(HWC3(edge)) + + return detected_map + +def safe_step(x, step=2): + y = x.astype(np.float32) * float(step + 1) + y = y.astype(np.int32).astype(np.float32) / float(step) + return y +def lineart_standard_preprocessor(image, guassian_sigma=6.0, intensity_threshold=8, resolution=1024): + """ + Standard lineart + :param image: Input image (should be a single image tensor) + :param guassian_sigma: Gaussian blur sigma + :param intensity_threshold: Threshold for intensity to define edges + :param resolution: Resolution for detection + :return: Processed image tensor + """ + model = LineartStandardDetector() + kwargs = {'guassian_sigma': guassian_sigma, 'intensity_threshold': intensity_threshold, 'resolution': resolution} + + detect_resolution = kwargs.pop('resolution', 512) if isinstance(kwargs.get('resolution', 512), int) and kwargs.get('resolution', 512) >= 64 else 512 + + # Convert tensor image to numpy array, apply model, and convert back to tensor + np_image = np.asarray(image.cpu() * 255., dtype=np.uint8) # Assuming the input image tensor is in the range [0, 1] + np_result = model(np_image, output_type="np", detect_resolution=detect_resolution, **kwargs) + tensor_result = torch.from_numpy(np_result.astype(np.float32) / 255.0) # Convert back to tensor in the range [0, 1] + + return tensor_result + +class LineartStandardDetector: + def __call__(self, input_image=None, guassian_sigma=6.0, intensity_threshold=8, detect_resolution=512, + output_type=None, upscale_method="INTER_CUBIC", **kwargs): + input_image, output_type = common_input_validate(input_image, output_type, **kwargs) + input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method) + + x = input_image.astype(np.float32) + g = cv2.GaussianBlur(x, (0, 0), guassian_sigma) + intensity = np.min(g - x, axis=2).clip(0, 255) + intensity /= max(16, np.median(intensity[intensity > intensity_threshold])) + intensity *= 127 + detected_map = intensity.clip(0, 255).astype(np.uint8) + + detected_map = HWC3(remove_pad(detected_map)) + + return detected_map +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 common_input_validate(input_image, output_type, **kwargs): + if "img" in kwargs: + # warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning) + input_image = kwargs.pop("img") + + if "return_pil" in kwargs: + # warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) + output_type = "pil" if kwargs["return_pil"] else "np" + + if type(output_type) is bool: + # warnings.warn( "Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") + if output_type: + output_type = "pil" + + if input_image is None: + raise ValueError("input_image must be defined.") + + if not isinstance(input_image, np.ndarray): + input_image = np.array(input_image, dtype=np.uint8) + output_type = output_type or "pil" + else: + output_type = output_type or "np" + + return (input_image, output_type) + +def resize_image_with_pad(input_image, resolution, upscale_method = "", skip_hwc3=False): + if skip_hwc3: + img = input_image + else: + img = HWC3(input_image) + H_raw, W_raw, _ = img.shape + k = float(resolution) / float(min(H_raw, W_raw)) + H_target = int(np.round(float(H_raw) * k)) + W_target = int(np.round(float(W_raw) * k)) + img = cv2.resize(img, (W_target, H_target), interpolation=get_upscale_method(upscale_method) if k > 1 else cv2.INTER_AREA) + H_pad, W_pad = pad64(H_target), pad64(W_target) + img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode='edge') + + def remove_pad(x): + return safer_memory(x[:H_target, :W_target, ...]) + + return safer_memory(img_padded), remove_pad +UPSCALE_METHODS = ["INTER_NEAREST", "INTER_LINEAR", "INTER_AREA", "INTER_CUBIC", "INTER_LANCZOS4"] +def get_upscale_method(method_str): + assert method_str in UPSCALE_METHODS, f"Method {method_str} not found in {UPSCALE_METHODS}" + return getattr(cv2, method_str) + +def pad64(x): + return int(np.ceil(float(x) / 64.0) * 64 - x) + +def safer_memory(x): + # Fix many MAC/AMD problems + return np.ascontiguousarray(x.copy()).copy() +def get_color_range(numpy_0_1, grey_num_min_0_1, grey_num_max_0_1, other_color=0, mode=1): + # 取灰度在某一范围内的像素grey_num_min__0_1,grey_num_max_0_1为灰度最大值最小值（包括等于）注意误差！！，other_color是所有不符合像素设成的颜色，mode=1时所有符合像素设成1，mode=2时所有符合像素保持不变 + result = numpy_0_1[:, :, 0] + if mode == 1: + result = np.where((result >= grey_num_min_0_1) & (result <= grey_num_max_0_1), 1, other_color) + elif mode == 2: + result = np.where((result >= grey_num_min_0_1) & (result <= grey_num_max_0_1), result, other_color) + result = np.expand_dims(result, 2).repeat(3, axis=2) + return result +def process_line_art(img,threshold=0.15,gaussian_sigma=2,intensity_threshold=3,min_size=36): + # threshold=0.5,gaussian_sigma=3,intensity_threshold=15,min_size=81 + lineart_standard_res = lineart_standard_preprocessor(image=img,guassian_sigma=gaussian_sigma, intensity_threshold=intensity_threshold,resolution=img.shape[0]).squeeze().numpy() + print("lineart_standard_res: ", lineart_standard_res) + lineart_standard_res = get_color_range(lineart_standard_res,threshold,1,other_color=0,mode=2) + print("lineart_standard_res 2: ", lineart_standard_res) + cleaned = morphology.remove_small_objects(lineart_standard_res.astype(bool), min_size=min_size, connectivity=1) + return lineart_standard_res*cleaned diff --git a/checkpoints/MTEED.pth b/checkpoints/MTEED.pth new file mode 100644 index 0000000..4fc5380 Binary files /dev/null and b/checkpoints/MTEED.pth differ diff --git a/requirements.txt b/requirements.txt index 84c49ea..000a1a1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ numpy kornia huggingface_hub opencv-python -scikit-image \ No newline at end of file +scikit-image +einops \ No newline at end of file diff --git a/teed_module.py b/teed_module.py index 9c1872e..e577e0d 100644 --- a/teed_module.py +++ b/teed_module.py @@ -3,10 +3,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from skimage import morphology -import numpy as np -import cv2 -from PIL import Image, ImageOps @torch.jit.script def mish_(input): @@ -207,146 +203,6 @@ class DoubleConvBlock(nn.Module): if self.use_act: x = self.smish(x) return x -def process_line_art(img, threshold=0.15, gaussian_sigma=2, intensity_threshold=3, min_size=36): - lineart_standard_res = lineart_standard_preprocessor(image=img, guassian_sigma=gaussian_sigma, intensity_threshold=intensity_threshold, resolution=img.shape[0]).squeeze().numpy() - lineart_standard_res = get_color_range(lineart_standard_res, threshold, 1, other_color=0, mode=2) - cleaned = morphology.remove_small_objects(lineart_standard_res.astype(bool), min_size=min_size, connectivity=1) - return lineart_standard_res * cleaned - -def lineart_standard_preprocessor(image, guassian_sigma=6.0, intensity_threshold=8, resolution=1024): - """ - 标准lineart - :param image: - :param guassian_sigma: - :param intensity_threshold: - :param resolution: - :return: - """ - model=LineartStandardDetector() - tensor_image=image - input_batch = False - kwargs= {'guassian_sigma':guassian_sigma, 'intensity_threshold':intensity_threshold, 'resolution':resolution} - if "detect_resolution" in kwargs: - del kwargs["detect_resolution"] #Prevent weird case? - - if "resolution" in kwargs: - detect_resolution = kwargs["resolution"] if type(kwargs["resolution"]) == int and kwargs["resolution"] >= 64 else 512 - del kwargs["resolution"] - else: - detect_resolution = 512 - - if input_batch: - np_images = np.asarray(tensor_image * 255., dtype=np.uint8) - np_results = model(np_images, output_type="np", detect_resolution=detect_resolution, **kwargs) - return torch.from_numpy(np_results.astype(np.float32) / 255.0) - - out_list = [] - for image in tensor_image: - np_image = np.asarray(image.cpu() * 255., dtype=np.uint8) - np_result = model(np_image, output_type="np", detect_resolution=detect_resolution, **kwargs) - out_list.append(torch.from_numpy(np_result.astype(np.float32) / 255.0)) - return torch.stack(out_list, dim=0) - -class LineartStandardDetector: - def __call__(self, input_image=None, guassian_sigma=6.0, intensity_threshold=8, detect_resolution=512, - output_type=None, upscale_method="INTER_CUBIC", **kwargs): - input_image, output_type = common_input_validate(input_image, output_type, **kwargs) - input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method) - - x = input_image.astype(np.float32) - g = cv2.GaussianBlur(x, (0, 0), guassian_sigma) - intensity = np.min(g - x, axis=2).clip(0, 255) - intensity /= max(16, np.median(intensity[intensity > intensity_threshold])) - intensity *= 127 - detected_map = intensity.clip(0, 255).astype(np.uint8) - - detected_map = HWC3(remove_pad(detected_map)) - if output_type == "pil": - detected_map = Image.fromarray(detected_map) - return detected_map - -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 common_input_validate(input_image, output_type, **kwargs): - if "img" in kwargs: - # warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning) - input_image = kwargs.pop("img") - - if "return_pil" in kwargs: - # warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) - output_type = "pil" if kwargs["return_pil"] else "np" - - if type(output_type) is bool: - # warnings.warn( "Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") - if output_type: - output_type = "pil" - - if input_image is None: - raise ValueError("input_image must be defined.") - - if not isinstance(input_image, np.ndarray): - input_image = np.array(input_image, dtype=np.uint8) - output_type = output_type or "pil" - else: - output_type = output_type or "np" - - return (input_image, output_type) - -def resize_image_with_pad(input_image, resolution, upscale_method = "", skip_hwc3=False): - if skip_hwc3: - img = input_image - else: - img = HWC3(input_image) - H_raw, W_raw, _ = img.shape - k = float(resolution) / float(min(H_raw, W_raw)) - H_target = int(np.round(float(H_raw) * k)) - W_target = int(np.round(float(W_raw) * k)) - img = cv2.resize(img, (W_target, H_target), interpolation=get_upscale_method(upscale_method) if k > 1 else cv2.INTER_AREA) - H_pad, W_pad = pad64(H_target), pad64(W_target) - img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode='edge') - - def remove_pad(x): - return safer_memory(x[:H_target, :W_target, ...]) - - return safer_memory(img_padded), remove_pad -UPSCALE_METHODS = ["INTER_NEAREST", "INTER_LINEAR", "INTER_AREA", "INTER_CUBIC", "INTER_LANCZOS4"] -def get_upscale_method(method_str): - assert method_str in UPSCALE_METHODS, f"Method {method_str} not found in {UPSCALE_METHODS}" - return getattr(cv2, method_str) - -def pad64(x): - return int(np.ceil(float(x) / 64.0) * 64 - x) - -def safer_memory(x): - # Fix many MAC/AMD problems - return np.ascontiguousarray(x.copy()).copy() -def get_color_range(numpy_0_1, grey_num_min_0_1, grey_num_max_0_1, other_color=0, mode=1): - # 取灰度在某一范围内的像素grey_num_min__0_1,grey_num_max_0_1为灰度最大值最小值（包括等于）注意误差！！，other_color是所有不符合像素设成的颜色，mode=1时所有符合像素设成1，mode=2时所有符合像素保持不变 - result = numpy_0_1[:, :, 0] - if mode == 1: - result = np.where((result >= grey_num_min_0_1) & (result <= grey_num_max_0_1), 1, other_color) - elif mode == 2: - result = np.where((result >= grey_num_min_0_1) & (result <= grey_num_max_0_1), result, other_color) - result = np.expand_dims(result, 2).repeat(3, axis=2) - return result -def combine_images(original, lineart): - # Logic to combine original and lineart images - pass class MTEED(nn.Module): def __init__(self): @@ -390,8 +246,7 @@ class MTEED(nn.Module): return data - def forward(self, x, original=None, is_eval=False): - assert original is not None + def forward(self, x): assert x.ndim == 4, x.shape block_1 = self.block_1(x) block_1_side = self.side_1(block_1) @@ -407,10 +262,4 @@ class MTEED(nn.Module): block_cat = torch.cat(results, dim=1) block_cat = self.block_cat(block_cat) results.append(block_cat) - if is_eval: - # add post processing - line_art_results = process_line_art(original) - combined_results = combine_images(results[0], line_art_results) - return combined_results - else: - return results \ No newline at end of file + return results \ No newline at end of file