diff --git a/anyline.py b/anyline.py index b9af680..f529d5f 100644 --- a/anyline.py +++ b/anyline.py @@ -1,4 +1,3 @@ -import os import torch import numpy as np import kornia as kn @@ -7,7 +6,6 @@ from .teed_module import MTEED from huggingface_hub import hf_hub_download import cv2 from skimage import morphology -from PIL import Image from einops import rearrange class AnyLine: @@ -49,16 +47,18 @@ class AnyLine: def get_anyline(self, image): # Process the image with MTEED model + print(image[0]) mteed_result = self.mteed_detector(image[0], detect_resolution=image[0].shape[0]) print("mteed_result: ", mteed_result) # Process the image with the lineart standard preprocessor - lineart_result = process_line_art(image[0]) + print(image) + lineart_result = process_line_art(image,threshold=0,gaussian_sigma=1,intensity_threshold=15,min_size=0) print("lineart_result: ", lineart_result) - + print("lineart_result: ", lineart_result.shape) # Combine the results - - + cv2.imshow('lineart_result', (lineart_result * 255.).astype(np.uint8)) + cv2.waitKey(0) return (torch.tensor(lineart_result),) class MTEEDDetector: @@ -85,31 +85,6 @@ class AnyLine: 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, @@ -127,6 +102,43 @@ class LineartStandardDetector: detected_map = HWC3(remove_pad(detected_map)) 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): + """ + 标准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) def HWC3(x): assert x.dtype == np.uint8 if x.ndim == 2: @@ -209,8 +221,6 @@ def get_color_range(numpy_0_1, grey_num_min_0_1, grey_num_max_0_1, other_color=0 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 + return lineart_standard_res*cleaned \ No newline at end of file