File size: 2,690 Bytes
baa8e90 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
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,) |