add pre-commit hooks configuration

This commit is contained in:
Tran Xen
2023-07-28 18:25:28 +02:00
parent 8577d0186d
commit 5d4a29ff1e
33 changed files with 1674 additions and 820 deletions
+40 -29
View File
@@ -1,6 +1,6 @@
import io
from typing import Optional
from PIL import Image, ImageChops, ImageOps,ImageFilter
from PIL import Image, ImageChops, ImageOps, ImageFilter
import cv2
import numpy as np
from math import isqrt, ceil
@@ -10,6 +10,7 @@ from scripts.faceswaplab_globals import NSFW_SCORE
from modules import processing
import base64
def check_against_nsfw(img):
shapes = []
chunks = detect(img)
@@ -17,6 +18,7 @@ def check_against_nsfw(img):
shapes.append(chunk["score"] > NSFW_SCORE)
return any(shapes)
def pil_to_cv2(pil_img):
return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
@@ -24,6 +26,7 @@ def pil_to_cv2(pil_img):
def cv2_to_pil(cv2_img):
return Image.fromarray(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB))
def torch_to_pil(images):
"""
Convert a numpy image or a batch of images to a PIL image.
@@ -49,35 +52,38 @@ def pil_to_torch(pil_images):
torch_image = torch.from_numpy(numpy_image).permute(2, 0, 1)
return torch_image
from collections import Counter
def create_square_image(image_list):
"""
Creates a square image by combining multiple images in a grid pattern.
Args:
image_list (list): List of PIL Image objects to be combined.
Returns:
PIL Image object: The resulting square image.
None: If the image_list is empty or contains only one image.
"""
# Count the occurrences of each image size in the image_list
size_counter = Counter(image.size for image in image_list)
# Get the most common image size (size with the highest count)
common_size = size_counter.most_common(1)[0][0]
# Filter the image_list to include only images with the common size
image_list = [image for image in image_list if image.size == common_size]
# Get the dimensions (width and height) of the common size
size = common_size
# If there are more than one image in the image_list
if len(image_list) > 1:
num_images = len(image_list)
# Calculate the number of rows and columns for the grid
rows = isqrt(num_images)
cols = ceil(num_images / rows)
@@ -97,10 +103,11 @@ def create_square_image(image_list):
# Return the resulting square image
return square_image
# Return None if there are no images or only one image in the image_list
return None
def create_mask(image, box_coords):
width, height = image.size
mask = Image.new("L", (width, height), 255)
@@ -113,43 +120,47 @@ def create_mask(image, box_coords):
mask.putpixel((x, y), 0)
return mask
def apply_mask(img : Image.Image,p : processing.StableDiffusionProcessing, batch_index : int) -> Image.Image :
def apply_mask(
img: Image.Image, p: processing.StableDiffusionProcessing, batch_index: int
) -> Image.Image:
"""
Apply mask overlay and color correction to an image if enabled
Args:
img: PIL Image objects.
p : The processing object
batch_index : the batch index
Returns:
PIL Image object
"""
if isinstance(p, processing.StableDiffusionProcessingImg2Img) :
if p.inpaint_full_res :
if isinstance(p, processing.StableDiffusionProcessingImg2Img):
if p.inpaint_full_res:
overlays = p.overlay_images
if overlays is None or batch_index >= len(overlays):
return img
overlay : Image.Image = overlays[batch_index]
overlay = overlay.resize((img.size), resample= Image.Resampling.LANCZOS)
overlay: Image.Image = overlays[batch_index]
overlay = overlay.resize((img.size), resample=Image.Resampling.LANCZOS)
img = img.copy()
img.paste(overlay, (0, 0), overlay)
return img
return img
img = processing.apply_overlay(img, p.paste_to, batch_index, p.overlay_images)
if p.color_corrections is not None and batch_index < len(p.color_corrections):
img = processing.apply_color_correction(p.color_corrections[batch_index], img)
img = processing.apply_color_correction(
p.color_corrections[batch_index], img
)
return img
def prepare_mask(
mask: Image.Image, p: processing.StableDiffusionProcessing
) -> Image.Image:
"""
Prepare an image mask for the inpainting process. (This comes from controlnet)
This function takes as input a PIL Image object and an instance of the
This function takes as input a PIL Image object and an instance of the
StableDiffusionProcessing class, and performs the following steps to prepare the mask:
1. Convert the mask to grayscale (mode "L").
@@ -160,26 +171,26 @@ def prepare_mask(
Args:
mask (Image.Image): The input mask as a PIL Image object.
p (processing.StableDiffusionProcessing): An instance of the StableDiffusionProcessing class
p (processing.StableDiffusionProcessing): An instance of the StableDiffusionProcessing class
containing the processing parameters.
Returns:
mask (Image.Image): The prepared mask as a PIL Image object.
"""
mask = mask.convert("L")
#FIXME : Properly fix blur
# FIXME : Properly fix blur
# if getattr(p, "mask_blur", 0) > 0:
# mask = mask.filter(ImageFilter.GaussianBlur(p.mask_blur))
return mask
def base64_to_pil(base64str : Optional[str]) -> Optional[Image.Image] :
if base64str is None :
def base64_to_pil(base64str: Optional[str]) -> Optional[Image.Image]:
if base64str is None:
return None
if 'base64,' in base64str: # check if the base64 string has a data URL scheme
base64_data = base64str.split('base64,')[-1]
if "base64," in base64str: # check if the base64 string has a data URL scheme
base64_data = base64str.split("base64,")[-1]
img_bytes = base64.b64decode(base64_data)
else:
# if no data URL scheme, just decode
img_bytes = base64.b64decode(base64str)
return Image.open(io.BytesIO(img_bytes))