Added High freq blending code

app.py

@@ -7,8 +7,16 @@ import time
 import logging
 import dotenv
 import fal_client
-import requests
 import base64
+import numpy as np
+import math
+import scipy
+import torch
+import torchvision
+import gradio as gr
+import argparse
+import spaces
+from PIL import Image, ImageFilter, ImageOps, ImageDraw, ImageFont
 from io import BytesIO
 from typing import Dict, List, Tuple, Union, Optional
 
@@ -34,14 +42,6 @@ sys.path.append(os.path.join(os.getcwd(), "GroundingDINO"))
 sys.path.append(os.path.join(os.getcwd(), "sam-hq"))
 warnings.filterwarnings("ignore")
 
-import numpy as np
-import torch
-import torchvision
-import gradio as gr
-import argparse
-from PIL import Image, ImageFilter, ImageOps, ImageDraw, ImageFont
-
-
 # Grounding DINO
 import GroundingDINO.groundingdino.datasets.transforms as T
 from GroundingDINO.groundingdino.models import build_model
@@ -56,10 +56,10 @@ CONFIG_FILE = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'
 GROUNDINGDINO_CHECKPOINT = "groundingdino_swint_ogc.pth"
 SAM_CHECKPOINT = 'sam_hq_vit_l.pth'
 OUTPUT_DIR = "outputs"
-FAL_KEY = os.getenv("FAL_KEY")
-UPLOAD_DIR = "./tmp/images"
+# FAL_KEY = os.getenv("FAL_KEY")
+# UPLOAD_DIR = "./tmp/images"
 
-os.makedirs(UPLOAD_DIR, exist_ok=True)
+# os.makedirs(UPLOAD_DIR, exist_ok=True)
 
 # Global variables for model caching
 _models = {
@@ -329,89 +329,106 @@ def run_grounded_sam(input_image):
         else:
             return [Image.new('RGB', (400, 300), color='gray'), Image.new('RGBA', (400, 300), color=(0, 0, 0, 0))]
 
-def image_gaussian_blur(image: torch.Tensor, radius: float) -> torch.Tensor:
-    if image.ndim == 4:  # Remove batch dimension if present
-        image = image.squeeze(0)
-    pil_image = tensor2pil(image)
-    blurred_pil_image = pil_image.filter(ImageFilter.GaussianBlur(radius))
-    return pil2tensor(blurred_pil_image).squeeze(0)
-
-def load_image(image_path: str) -> torch.Tensor:
-    image = Image.open(image_path).convert("RGBA")
-    image_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0
-    return image_tensor
-
-def split_image_with_alpha(image: torch.Tensor):
-    out_images = image[:3, :, :]
-    out_alphas = image[3, :, :] if image.shape[0] > 3 else torch.ones_like(image[0, :, :])
-    result = (out_images.unsqueeze(0), 1.0 - out_alphas.unsqueeze(0))
-    return result
-
-def pil2numpy(image: Image.Image):
-    return np.array(image).astype(np.float32) / 255.0
-
-def numpy2pil(image: np.ndarray, mode=None):
-    return Image.fromarray(np.clip(255.0 * image, 0, 255).astype(np.uint8), mode)
+def split_image_with_alpha(image):
+    image = image.convert("RGB")
+    return image
 
-def pil2tensor(image: Image.Image):
-    return torch.from_numpy(pil2numpy(image)).unsqueeze(0)
+def gaussian_blur(image, radius=10):
+    """Apply Gaussian blur to image."""
+    blurred = image.filter(ImageFilter.GaussianBlur(radius=10))
+    return blurred
 
-def invert(image):
-    s = 1.0 - image
-    return s
+def invert_image(image):
+    img_inverted = ImageOps.invert(image)
+    return img_inverted
 
-def image2tensor(image) -> torch.Tensor:
-    image_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0
-    return image_tensor
+def expand_mask(mask, expand, tapered_corners):
+    # Ensure mask is in grayscale (mode 'L')
+    mask = mask.convert("L")
+    
+    # Convert to NumPy array
+    mask_np = np.array(mask)
 
-def tensor2pil(image: torch.Tensor, mode=None):
-    if image.ndim == 2:  # Grayscale image
-        image = image.unsqueeze(0)  # Add channel dimension
+    # Define kernel
+    c = 0 if tapered_corners else 1
+    kernel = np.array([[c, 1, c],
+                       [1, 1, 1],
+                       [c, 1, c]], dtype=np.uint8)
 
-    if image.ndim != 3 or image.shape[1:] == (0, 0):
-        raise ValueError(f"Invalid tensor dimensions: {image.shape}")
+    # Perform dilation or erosion based on expand value
+    if expand > 0:
+        for _ in range(expand):
+            mask_np = scipy.ndimage.grey_dilation(mask_np, footprint=kernel)
+    elif expand < 0:
+        for _ in range(abs(expand)):
+            mask_np = scipy.ndimage.grey_erosion(mask_np, footprint=kernel)
 
-    if image.shape[0] == 1:  # Single channel, replicate to 3 channels
-        image = image.repeat(3, 1, 1)
-    elif image.shape[0] != 3:
-        raise ValueError("Unexpected number of channels in the image tensor")
-    
-    return numpy2pil(image.cpu().numpy().transpose(1, 2, 0), mode=mode)
+    # Convert back to PIL image
+    return Image.fromarray(mask_np, mode="L")
 
-def extract_high_frequency(image: torch.Tensor, blur_radius: float = 5.0) -> torch.Tensor:
-    """Extract high-frequency details by subtracting the blurred image from the original."""
-    if image.ndim == 4:
-        image = image.squeeze(0)
-    
-    blurred = image_gaussian_blur(image, blur_radius)
-    
-    if blurred.ndim == 4:
-        blurred = blurred.squeeze(0)
-    elif blurred.ndim == 3 and blurred.shape[0] != 3:
-        blurred = blurred.permute(2, 0, 1)
-    
-    high_freq = image - blurred
-    return high_freq
+def image_blend_by_mask(image_a, image_b, mask, blend_percentage):
 
-def image_blend_mask(image_a, image_b, mask, blend_percentage):
-
-    # Convert images to PIL
-    img_a = tensor2pil(image_a)
-    img_b = tensor2pil(image_b)
-    mask = ImageOps.invert(tensor2pil(mask).convert('L'))
+    mask = ImageOps.invert(mask.convert('L'))
 
     # Mask image
-    masked_img = Image.composite(img_a, img_b, mask.resize(img_a.size))
+    masked_img = Image.composite(image_a, image_b, mask)
 
     # Blend image
-    blend_mask = Image.new(mode="L", size=img_a.size,
+    blend_mask = Image.new(mode="L", size=image_a.size,
                            color=(round(blend_percentage * 255)))
     blend_mask = ImageOps.invert(blend_mask)
-    img_result = Image.composite(img_a, masked_img, blend_mask)
+    img_result = Image.composite(image_a, masked_img, blend_mask)
+
+    del image_a, image_b, blend_mask, mask
+
+    return img_result
+
+def blend_images(image_a, image_b, blend_percentage):
+    """Blend img_b over image_a using the normal mode with a blend percentage."""
+    img_a = image_a.convert("RGBA")
+    img_b = image_b.convert("RGBA")
+
+    # Blend img_b over img_a using alpha_composite (normal blend mode)
+    out_image = Image.alpha_composite(img_a, img_b)
+
+    out_image = out_image.convert("RGB")
+
+    # Create blend mask
+    blend_mask = Image.new("L", image_a.size, round(blend_percentage * 255))
+    blend_mask = ImageOps.invert(blend_mask)  # Invert the mask
 
-    del img_a, img_b, blend_mask, mask
+    # Apply composite blend
+    result = Image.composite(image_a, out_image, blend_mask)
+    return result
+
+def apply_image_levels(image, black_level, mid_level, white_level):
+    levels = AdjustLevels(black_level, mid_level, white_level)
+    adjusted_image = levels.adjust(image)
+    return adjusted_image
+
+class AdjustLevels:
+    def __init__(self, min_level, mid_level, max_level):
+        self.min_level = min_level
+        self.mid_level = mid_level
+        self.max_level = max_level
+
+    def adjust(self, im):
+
+        im_arr = np.array(im).astype(np.float32)
+        im_arr[im_arr < self.min_level] = self.min_level
+        im_arr = (im_arr - self.min_level) * \
+            (255 / (self.max_level - self.min_level))
+        im_arr = np.clip(im_arr, 0, 255)
+
+        # mid-level adjustment
+        gamma = math.log(0.5) / math.log((self.mid_level - self.min_level) / (self.max_level - self.min_level))
+        im_arr = np.power(im_arr / 255, gamma) * 255
+
+        im_arr = im_arr.astype(np.uint8)
+
+        im = Image.fromarray(im_arr)
 
-    return (pil2tensor(img_result), )
+        return im
 
 def encode_image(image):
     buffer = BytesIO()
@@ -442,75 +459,71 @@ def generate_ai_bg(input_img, prompt):
     return ic_light_img
 
 def blend_details(input_image, relit_image, masked_image):
-    with torch.inference_mode():
-        # Convert images to tensors
-        input_image = image2tensor(input_image)
-        relit_image = image2tensor(relit_image)
-        masked_image = image2tensor(masked_image)
-
-        # Resize input image
-        input_image = torch.nn.functional.interpolate(
-            input_image.unsqueeze(0),
-            size=(1024, 1024),
-            mode="bicubic",
-            align_corners=False
-        ).squeeze(0)
-
-        # Resize relit image
-        relit_image = torch.nn.functional.interpolate(
-            relit_image.unsqueeze(0),
-            size=(1024, 1024),
-            mode="bicubic",
-            align_corners=False
-        ).squeeze(0)
-
-        # Resize masked image
-        masked_image = torch.nn.functional.interpolate(
-            masked_image.unsqueeze(0),
-            size=(1024, 1024),
-            mode="bicubic",
-            align_corners=False
-        ).squeeze(0)
-
-        # Split images and get RGB channels
-        input_image_rgb = split_image_with_alpha(input_image)[0].squeeze(0)
-        relit_image_rgb = split_image_with_alpha(relit_image)[0].squeeze(0)
-        
-        # Use masked image RGB channels as segmentation mask (average of RGB channels)
-        segmentation_mask = masked_image[:3].mean(dim=0)  # Average RGB channels to get grayscale mask
-        
-        print(f"segmentation_mask shape: {segmentation_mask.shape}")
-        
-        # Extract high-frequency details from input image
-        high_freq_details = extract_high_frequency(input_image_rgb, blur_radius=3.0)
-
-        # Print shapes for debugging
-        print(f"high_freq_details shape: {high_freq_details.shape}")
-        print(f"segmentation_mask shape: {segmentation_mask.shape}")
-        print(f"relit_image_rgb shape: {relit_image_rgb.shape}")
-
-        # Apply high-frequency details only in masked areas
-        detail_strength = 0.5
-        segmentation_mask = segmentation_mask.unsqueeze(0).repeat(3, 1, 1)  # Expand mask to match RGB channels
-        masked_details = high_freq_details * segmentation_mask
-        # final_image = relit_image_rgb + (masked_details * detail_strength)
-        # final_image = image_blend_mask(relit_image_rgb, masked_details, mask, blend_percentage)
-        final_image = relit_image_rgb + masked_details
-        print('final_image shape:', final_image.shape)
-
-        # Normalize to [0, 1] range
-        final_image = torch.clamp(final_image, 0, 1)
-
-        # Save intermediate results for debugging
-        # tensor2pil(segmentation_mask).save("output/segmentation_mask.png")
-        # tensor2pil(high_freq_details).save("output/high_freq_details.png")
-        # tensor2pil(masked_details).save("output/masked_details.png")
+    # input_image = load_image(input_image_path)
+    # relit_image = load_image(relit_image_path)
+    # masked_image = load_image(masked_image_path)
+
+    scaling_factor = 1
+    input_image = input_image.resize((int(input_image.width * scaling_factor), 
+        int(input_image.height * scaling_factor)))
+
+    relit_image = relit_image.resize((int(relit_image.width * scaling_factor), 
+        int(relit_image.height * scaling_factor)))
+
+    masked_image = masked_image.resize((int(masked_image.width * scaling_factor), 
+        int(masked_image.height * scaling_factor)))
+
+    masked_image_rgb = split_image_with_alpha(masked_image)
+    masked_image_blurred = gaussian_blur(masked_image_rgb, radius=10)
+    grow_mask = expand_mask(masked_image_blurred, -15, True)
+
+    # grow_mask.save("output/grow_mask.png")
+
+    # Split images and get RGB channels
+    input_image_rgb = split_image_with_alpha(input_image)
+    input_blurred = gaussian_blur(input_image_rgb, radius=10)
+    input_inverted = invert_image(input_image_rgb)
+
+    # input_blurred.save("output/input_blurred.png")
+    # input_inverted.save("output/input_inverted.png")
         
-        # Save final result
-        final_image_pil = tensor2pil(final_image)
-        # final_image_pil.save("output/output_image.png")
-        return final_image_pil
+    # Add blurred and inverted images
+    input_blend_1 = blend_images(input_inverted, input_blurred, blend_percentage=0.5)
+    input_blend_1_inverted = invert_image(input_blend_1)
+    input_blend_2 = blend_images(input_blurred, input_blend_1_inverted, blend_percentage=1.0)
+
+    # input_blend_2.save("output/input_blend_2.png")
+
+    # Process relit image
+    relit_image_rgb = split_image_with_alpha(relit_image)
+    relit_blurred = gaussian_blur(relit_image_rgb, radius=10)
+    relit_inverted = invert_image(relit_image_rgb)
+
+    # relit_blurred.save("output/relit_blurred.png")
+    # relit_inverted.save("output/relit_inverted.png")
+    
+    # Add blurred and inverted relit images
+    relit_blend_1 = blend_images(relit_inverted, relit_blurred, blend_percentage=0.5)
+    relit_blend_1_inverted = invert_image(relit_blend_1)
+    relit_blend_2 = blend_images(relit_blurred, relit_blend_1_inverted, blend_percentage=1.0)
+
+    # relit_blend_2.save("output/relit_blend_2.png")
+
+    high_freq_comp = image_blend_by_mask(relit_blend_2, input_blend_2, grow_mask, blend_percentage=1.0)
+
+    # high_freq_comp.save("output/high_freq_comp.png")
+
+    comped_image = blend_images(relit_blurred, high_freq_comp, blend_percentage=0.65)
+
+    # comped_image.save("output/comped_image.png")
+
+    final_image = apply_image_levels(comped_image, black_level=83, mid_level=128, white_level=172)
+
+    # final_image.save("output/final_image.png")
+
+    return final_image
 
+@spaces.GPU
 def generate_image(input_img, ai_gen_image, prompt):
 
     # ai_gen_image = generate_ai_bg(input_img, prompt)
@@ -539,7 +552,7 @@ def create_ui():
                 # gallery = gr.Gallery(
                 #     label="Generated images", show_label=False, elem_id="gallery"
                 # )
-                masked_image = gr.Image(label="Generated Image")
+                # masked_image = gr.Image(label="Generated Image")
                 output_image = gr.Image(label="Generated Image")
         
         # Run button
@@ -550,7 +563,7 @@ def create_ui():
                 ai_image,
                 prompt
             ], 
-            outputs=[masked_image, output_image]
+            outputs=[output_image]
         )     
  
     return block