Update src/utils.py

src/utils.py

@@ -1,10 +1,11 @@
+# utils.py
 from enum import Enum, auto
 
 import torch
-from huggingface_hub import (  # pyright: ignore[reportMissingTypeStubs]
-    hf_hub_download,  # pyright: ignore[reportUnknownVariableType]
-)
-from PIL import Image
+from huggingface_hub import hf_hub_download
+from PIL import Image, ImageEnhance, ImageFilter
+import cv2
+import numpy as np
 from refiners.fluxion.utils import load_from_safetensors, tensor_to_image
 from refiners.foundationals.clip import CLIPTextEncoderL
 from refiners.foundationals.latent_diffusion import SD1UNet
@@ -53,10 +54,7 @@ def resize_modulo_8(
     resample: Image.Resampling | None = None,
     on_short: bool = True,
 ) -> Image.Image:
-    """Resize an image respecting the aspect ratio and ensuring the size is a multiple of 8.
-
-    The `on_short` parameter determines whether the resizing is based on the shortest side.
-    """
+    """이미지 크기를 8의 배수로 조정"""
     assert size % 8 == 0, "Size must be a multiple of 8 because this is the latent compression size."
     side_size = min(image.size) if on_short else max(image.size)
     scale = size / (side_size * 8)
@@ -64,6 +62,95 @@ def resize_modulo_8(
     return image.resize(new_size, resample=resample or Image.Resampling.LANCZOS)
 
 
+def adjust_image(
+    image: Image.Image,
+    brightness=0.0,
+    contrast=0.0,
+    temperature=0.0,
+    saturation=0.0,
+    tint=0.0,
+    blur_intensity=0,
+    exposure=0.0,
+    vibrance=0.0,
+    color_mixer_blues=0.0,
+) -> Image.Image:
+    """이미지 조정 함수"""
+    image = image.convert('RGB')
+
+    # 노출 조정 (Exposure)
+    if exposure != 0.0:
+        # Exposure ranges from -5 to 5, where 0 is neutral
+        exposure_factor = 1 + (exposure / 5.0)
+        exposure_factor = max(exposure_factor, 0.01)  # Prevent zero or negative
+        enhancer = ImageEnhance.Brightness(image)
+        image = enhancer.enhance(exposure_factor)
+
+    # 밝기 조정
+    if brightness != 0.0:
+        # Brightness ranges from -5 to 5, mapped to brightness factor
+        brightness_factor = 1 + (brightness / 5.0)
+        brightness_factor = max(brightness_factor, 0.01)  # Prevent zero or negative
+        enhancer = ImageEnhance.Brightness(image)
+        image = enhancer.enhance(brightness_factor)
+
+    # 대비 조정
+    if contrast != 0.0:
+        # Contrast ranges from -100 to 100, mapped to contrast factor
+        contrast_factor = 1 + (contrast / 100.0)
+        contrast_factor = max(contrast_factor, 0.01)  # Prevent zero or negative
+        enhancer = ImageEnhance.Contrast(image)
+        image = enhancer.enhance(contrast_factor)
+
+    # 채도 조정 (Vibrance)
+    if vibrance != 0.0:
+        # Vibrance simulates adjusting the saturation; positive increases saturation, negative decreases
+        vibrance_factor = 1 + (vibrance / 100.0)
+        vibrance_factor = max(vibrance_factor, 0.0)  # Prevent negative saturation
+        enhancer = ImageEnhance.Color(image)
+        image = enhancer.enhance(vibrance_factor)
+
+    # 채도 조정 (Saturation)
+    if saturation != 0.0:
+        # Saturation ranges from -100 to 100, mapped to saturation factor
+        saturation_factor = 1 + (saturation / 100.0)
+        saturation_factor = max(saturation_factor, 0.0)  # Prevent negative saturation
+        enhancer = ImageEnhance.Color(image)
+        image = enhancer.enhance(saturation_factor)
+
+    # 색온도 조정
+    if temperature != 0.0:
+        # To prevent division by zero, adjust temperature calculation
+        temp_factor = 1 + (temperature / 100.0)
+        temp_factor = max(temp_factor, 0.01)  # Prevent zero or negative
+
+        r, g, b = image.split()
+        r = r.point(lambda i: i * temp_factor)
+        b = b.point(lambda i: i / temp_factor)
+        image = Image.merge('RGB', (r, g, b))
+
+    # 색조 조정 (Tint)
+    if tint != 0.0:
+        image_np = np.array(image)
+        image_hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV).astype(np.float32)
+        image_hsv[:, :, 0] = (image_hsv[:, :, 0] + tint) % 180
+        image_hsv[:, :, 0] = np.clip(image_hsv[:, :, 0], 0, 179)
+        image_rgb = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)
+        image = Image.fromarray(image_rgb)
+
+    # 블러 적용
+    if blur_intensity > 0:
+        image = image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
+
+    # Color Mixer (Blues)
+    if color_mixer_blues != 0.0:
+        image_np = np.array(image).astype(np.float32)
+        # Adjust the blue channel
+        image_np[:, :, 2] = np.clip(image_np[:, :, 2] + (color_mixer_blues / 100.0) * 255, 0, 255)
+        image = Image.fromarray(image_np.astype(np.uint8))
+
+    return image
+
+
 class LightingPreference(str, Enum):
     LEFT = auto()
     RIGHT = auto()
@@ -72,12 +159,7 @@ class LightingPreference(str, Enum):
     NONE = auto()
 
     def get_init_image(self, width: int, height: int, interval: tuple[float, float] = (0.0, 1.0)) -> Image.Image | None:
-        """Generate an image with a linear gradient based on the lighting preference.
-
-        In the original code, interval is always (0., 1.) ; we added it as a parameter to make the function more
-        flexible and allow for less contrasted images with a smaller interval.
-        see https://github.com/lllyasviel/IC-Light/blob/7886874/gradio_demo.py#L242
-        """
+        """조명 선호도에 따른 그라데이션 이미지 생성"""
         start, end = interval
         match self:
             case LightingPreference.LEFT:
@@ -107,4 +189,4 @@ class LightingPreference(str, Enum):
             case "none":
                 return LightingPreference.NONE
             case _:
-                raise ValueError(f"Invalid lighting preference: {value}")
+                raise ValueError(f"Invalid lighting preference: {value}")