| import torch |
| import torch.nn.functional as F |
| import os |
| import subprocess |
|
|
| CV2_AVAILABLE = True |
| try: |
| import cv2 |
| except: |
| print("OpenCV is not installed so face cropping is not available.") |
| CV2_AVAILABLE = False |
|
|
| CURRENT_DIR = os.path.dirname(os.path.realpath(__file__)) |
| DETECTOR_FILE = "lbpcascade_animeface.xml" |
|
|
| if not os.path.exists(os.path.join(CURRENT_DIR, DETECTOR_FILE)): |
| print("Downloading anime face detector...") |
| try: |
| subprocess.run(["wget", "https://raw.githubusercontent.com/nagadomi/lbpcascade_animeface/master/lbpcascade_animeface.xml", "-P", CURRENT_DIR]) |
| except: |
| print(f"Failed to download lbpcascade_animeface.xml so please download it in {CURRENT_DIR}.") |
| CV2_AVAILABLE = False |
|
|
| CROP_MODES = ["padding", "face_crop", "none"] if CV2_AVAILABLE else ["padding", "none"] |
|
|
| def image_to_numpy(image): |
| image = image.squeeze(0) * 255 |
| return image.numpy().astype("uint8") |
|
|
| def numpy_to_image(image): |
| image = torch.tensor(image).float() / 255 |
| return image.unsqueeze(0) |
|
|
| def pad_to_square(tensor): |
| tensor = tensor.squeeze(0).permute(2, 0, 1) |
| _, h, w = tensor.shape |
|
|
| target_length = max(h, w) |
|
|
| pad_l = (target_length - w) // 2 |
| pad_r = (target_length - w) - pad_l |
| |
| pad_t = (target_length - h) // 2 |
| pad_b = (target_length - h) - pad_t |
|
|
| padded_tensor = F.pad(tensor, (pad_l, pad_r, pad_t, pad_b), mode="constant", value=0) |
|
|
| return padded_tensor.permute(1, 2, 0).unsqueeze(0) |
|
|
| def face_crop(image): |
| image = image_to_numpy(image) |
| face_cascade = cv2.CascadeClassifier(os.path.join(CURRENT_DIR, DETECTOR_FILE)) |
| gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) |
| faces = face_cascade.detectMultiScale(gray, 1.3, 5) |
|
|
| w, h = image.shape[1], image.shape[0] |
|
|
| target_length = min(w, h) |
| fx, fy, fw, fh = (0, 0, w, h) if len(faces) == 0 else faces[0] |
|
|
| dx = target_length - fw // 2 |
| dy = target_length - fh // 2 |
|
|
| target_x = 0 if w < h else max(0, fx - dx) |
| target_y = 0 if w > h else max(0, fy - dy) |
| |
| image = image[target_y:target_y+target_length, target_x:target_x+target_length] |
| image = numpy_to_image(image) |
|
|
| return image |
|
|
| class ImageCrop: |
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "image": ("IMAGE", ), |
| "mode": (CROP_MODES, ), |
| } |
| } |
| |
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "preprocess" |
| CATEGORY = "image/preprocessors" |
|
|
| def preprocess(self, image, mode): |
| if mode == "padding": |
| image = pad_to_square(image) |
| elif mode == "face_crop": |
| image = face_crop(image) |
| |
| return (image,) |
