From 5f0f792b34c0e75b241a9472e0a97409395b257c Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 13:54:42 +0800 Subject: [PATCH 1/6] :sparkles: chore: Update dependencies and add create_mask_from_sketch function --- pyproject.toml | 1 + src/novelai_python/utils/useful.py | 50 ++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 856ef44..5b27f39 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "uvicorn[standard]>=0.27.0.post1", "numpy>=1.24.4", "argon2-cffi>=23.1.0", + "opencv-python>=4.8.1.78", ] requires-python = ">=3.8" readme = "README.md" diff --git a/src/novelai_python/utils/useful.py b/src/novelai_python/utils/useful.py index 5a4a7e2..14bcc19 100644 --- a/src/novelai_python/utils/useful.py +++ b/src/novelai_python/utils/useful.py @@ -6,6 +6,10 @@ import collections import random from typing import List, Union +from typing import Optional + +import cv2 as cv +import numpy as np def enum_to_list(enum_): @@ -42,3 +46,49 @@ def get(self, user_id: Union[int, str]) -> str: user_used.append(selected) return selected + + +def create_mask_from_sketch(original_img_bytes: bytes, + sketch_img_bytes: bytes, + output_format: Optional[str] = '.png' + ) -> bytes: + """ + Function to create a mask from original and sketch images input as bytes. Returns bytes. + + :param original_img_bytes: Bytes corresponding to the original image. + :param sketch_img_bytes: Bytes corresponding to the sketch image. + :param output_format: Format of the output image. Defaults to '.png'. It could also be '.jpg' + :returns bytes: Bytes corresponding to the resultant mask + """ + # Load images + ori_img = cv.imdecode(np.frombuffer(original_img_bytes, np.uint8), cv.IMREAD_COLOR) + sketch_img = cv.imdecode(np.frombuffer(sketch_img_bytes, np.uint8), cv.IMREAD_COLOR) + + # Check if images have the same size + if ori_img.shape != sketch_img.shape: + raise ValueError("Images must have the same size.") + + # Calculate difference between the original and sketch images + diff_img = cv.absdiff(ori_img, sketch_img) + + # Convert the difference image to grayscale + diff_gray = cv.cvtColor(diff_img, cv.COLOR_BGR2GRAY) + + # Threshold to create the mask + _, thresh = cv.threshold(diff_gray, 10, 255, cv.THRESH_BINARY) + + # Perform morphological opening to remove small noise + kernel = np.ones((3, 3), np.uint8) + opening = cv.morphologyEx(thresh, cv.MORPH_OPEN, kernel, iterations=2) + + # 高斯模糊 + smooth = cv.GaussianBlur(opening, (5, 5), 0) + + # Convert image to bytes + is_success, buffer = cv.imencode(output_format, smooth) + if is_success: + output_bytes = buffer.tobytes() + else: + raise ValueError("Error during conversion of image to bytes") + + return output_bytes From de7930bf168311d05d3f2765a721947edc93f242 Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 13:54:51 +0800 Subject: [PATCH 2/6] :sparkles: chore: Update dependencies and add create_mask_from_sketch function --- pdm.lock | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/pdm.lock b/pdm.lock index 637b037..d699a32 100644 --- a/pdm.lock +++ b/pdm.lock @@ -5,7 +5,7 @@ groups = ["default", "dev", "testing"] strategy = ["cross_platform", "inherit_metadata"] lock_version = "4.4.1" -content_hash = "sha256:601014f72651e21a1c2c5592a2c66fa00ce682f0d7fc1b9ddcfd698ad7eeba30" +content_hash = "sha256:9aa6958e83642a27ca1aae9b488f8f01f014246d5f42b49f2cf7a04fc6328cb1" [[package]] name = "annotated-types" @@ -576,6 +576,32 @@ files = [ {file = "numpy-1.24.4.tar.gz", hash = "sha256:80f5e3a4e498641401868df4208b74581206afbee7cf7b8329daae82676d9463"}, ] +[[package]] +name = "opencv-python" +version = "4.8.1.78" +requires_python = ">=3.6" +summary = "Wrapper package for OpenCV python bindings." +groups = ["default"] +dependencies = [ + "numpy>=1.17.0; python_version >= \"3.7\"", + "numpy>=1.17.3; python_version >= \"3.8\"", + "numpy>=1.19.3; python_version >= \"3.6\" and platform_system == \"Linux\" and platform_machine == \"aarch64\"", + "numpy>=1.19.3; python_version >= \"3.9\"", + "numpy>=1.21.0; python_version <= \"3.9\" and platform_system == \"Darwin\" and platform_machine == \"arm64\"", + "numpy>=1.21.2; python_version >= \"3.10\"", + "numpy>=1.21.4; python_version >= \"3.10\" and platform_system == \"Darwin\"", + "numpy>=1.23.5; python_version >= \"3.11\"", +] +files = [ + {file = "opencv-python-4.8.1.78.tar.gz", hash = "sha256:cc7adbbcd1112877a39274106cb2752e04984bc01a031162952e97450d6117f6"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-macosx_10_16_x86_64.whl", hash = "sha256:91d5f6f5209dc2635d496f6b8ca6573ecdad051a09e6b5de4c399b8e673c60da"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:bc31f47e05447da8b3089faa0a07ffe80e114c91ce0b171e6424f9badbd1c5cd"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9814beca408d3a0eca1bae7e3e5be68b07c17ecceb392b94170881216e09b319"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c4c406bdb41eb21ea51b4e90dfbc989c002786c3f601c236a99c59a54670a394"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-win32.whl", hash = "sha256:a7aac3900fbacf55b551e7b53626c3dad4c71ce85643645c43e91fcb19045e47"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-win_amd64.whl", hash = "sha256:b983197f97cfa6fcb74e1da1802c7497a6f94ed561aba6980f1f33123f904956"}, +] + [[package]] name = "packaging" version = "23.2" From fc418c140145d3b6ee91aa0a3cbee1949b449822 Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 13:55:06 +0800 Subject: [PATCH 3/6] :sparkles: feat(generate_image): Add endpoint property and setter --- src/novelai_python/sdk/ai/generate_image.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/novelai_python/sdk/ai/generate_image.py b/src/novelai_python/sdk/ai/generate_image.py index 231910d..2fff470 100644 --- a/src/novelai_python/sdk/ai/generate_image.py +++ b/src/novelai_python/sdk/ai/generate_image.py @@ -207,6 +207,14 @@ def height_validator(cls, v: int): parameters: Params = Params() model_config = ConfigDict(extra="ignore") + @property + def endpoint(self): + return self._endpoint + + @endpoint.setter + def endpoint(self, value): + self._endpoint = value + @override def model_post_init(self, *args) -> None: """ @@ -241,20 +249,14 @@ def model_post_init(self, *args) -> None: def validate_model(self): if self.action == Action.INFILL and not self.parameters.mask: logger.warning("Mask maybe required for infill mode.") + if self.action == Action.INFILL: + self.parameters.extra_noise_seed = self.parameters.seed return self @property def base_url(self): return f"{self.endpoint.strip('/')}/ai/generate-image" - @property - def endpoint(self): - return self._endpoint - - @endpoint.setter - def endpoint(self, value): - self._endpoint = value - def calculate_cost(self, is_opus: bool = False): """ Calculate the Anlas cost of current parameters. From 67a48b78b18ae4443cac04c514f0218baa4b78fb Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 13:58:28 +0800 Subject: [PATCH 4/6] :sparkles: feat: Add mask creation function Add a new function to create a mask from original and sketch images input as bytes. The function includes an option to create jagged edges in the mask. The output format of the mask image can be specified as well. --- playground/mask/__init__.py | 17 +++++++++++++++++ src/novelai_python/utils/useful.py | 26 ++++++++++++++++++-------- 2 files changed, 35 insertions(+), 8 deletions(-) create mode 100644 playground/mask/__init__.py diff --git a/playground/mask/__init__.py b/playground/mask/__init__.py new file mode 100644 index 0000000..afb3fe2 --- /dev/null +++ b/playground/mask/__init__.py @@ -0,0 +1,17 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/2/13 下午12:42 +# @Author : sudoskys +# @File : __init__.py.py +# @Software: PyCharm +from novelai_python.utils.useful import create_mask_from_sketch + +with open('sk.jpg', 'rb') as f: + sk_bytes = f.read() + +with open('ori.png', 'rb') as f: + ori_bytes = f.read() + +return_bytes = create_mask_from_sketch(original_img_bytes=ori_bytes, sketch_img_bytes=sk_bytes, jagged_edges=True) + +with open('mask_export.png', 'wb') as f: + f.write(return_bytes) diff --git a/src/novelai_python/utils/useful.py b/src/novelai_python/utils/useful.py index 14bcc19..db52059 100644 --- a/src/novelai_python/utils/useful.py +++ b/src/novelai_python/utils/useful.py @@ -6,7 +6,6 @@ import collections import random from typing import List, Union -from typing import Optional import cv2 as cv import numpy as np @@ -50,11 +49,13 @@ def get(self, user_id: Union[int, str]) -> str: def create_mask_from_sketch(original_img_bytes: bytes, sketch_img_bytes: bytes, - output_format: Optional[str] = '.png' + output_format: str = '.png', + jagged_edges: bool = True ) -> bytes: """ Function to create a mask from original and sketch images input as bytes. Returns bytes. + :param jagged_edges: Whether to create jagged edges in the mask. Defaults to False. :param original_img_bytes: Bytes corresponding to the original image. :param sketch_img_bytes: Bytes corresponding to the sketch image. :param output_format: Format of the output image. Defaults to '.png'. It could also be '.jpg' @@ -77,18 +78,27 @@ def create_mask_from_sketch(original_img_bytes: bytes, # Threshold to create the mask _, thresh = cv.threshold(diff_gray, 10, 255, cv.THRESH_BINARY) + if jagged_edges: + # Use a bigger kernel for dilation to create larger 'step' effect at the edges + kernel = np.ones((7, 7), np.uint8) + else: + kernel = np.ones((3, 3), np.uint8) + # Perform morphological opening to remove small noise - kernel = np.ones((3, 3), np.uint8) opening = cv.morphologyEx(thresh, cv.MORPH_OPEN, kernel, iterations=2) - # 高斯模糊 + # Further remove noise with a Gaussian filter smooth = cv.GaussianBlur(opening, (5, 5), 0) - # Convert image to bytes + if jagged_edges: + # Apply additional thresholding to create sharper, jagged edges + _, smooth = cv.threshold(smooth, 128, 255, cv.THRESH_BINARY) + + # Convert image to BytesIO object is_success, buffer = cv.imencode(output_format, smooth) if is_success: - output_bytes = buffer.tobytes() + output_io = buffer.tobytes() else: - raise ValueError("Error during conversion of image to bytes") + raise ValueError("Error during conversion of image to BytesIO object") - return output_bytes + return output_io From 057ddd2176587c45f8e7fb4624e5b0ef66dc2cf7 Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 14:00:50 +0800 Subject: [PATCH 5/6] :sparkles: feat(playground): add enhance.py file --- playground/enhance.py | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 playground/enhance.py diff --git a/playground/enhance.py b/playground/enhance.py new file mode 100644 index 0000000..9c3b2c4 --- /dev/null +++ b/playground/enhance.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/2/13 下午1:58 +# @Author : sudoskys +# @File : enhance.py +# @Software: PyCharm + +# NOTE About Enhance Mode +# Enhance Mode = origin model name + img2img action + width *1.5(or 1) +height *1.5(or 1) + change seed +# :) From be9d2a12716255fcb9a20152300caf3ae752550b Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 14:01:55 +0800 Subject: [PATCH 6/6] :arrow_up: chore: bump version to 0.3.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 5b27f39..c3af90f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "novelai-python" -version = "0.2.9" +version = "0.3.0" description = "Novelai Python Binding With Pydantic" authors = [ { name = "sudoskys", email = "coldlando@hotmail.com" },