Update app.py

app.py

@@ -1,3 +1,224 @@
+'''
+import os
+from datasets import load_dataset
+from PIL import Image
+
+def save_images_locally(dataset_name="svjack/InfiniteYou_PosterCraft_Wang_Leehom_Poster_FP8_WAV_text_mask"):
+    # 加载数据集
+    dataset = load_dataset(dataset_name)["train"]
+
+    # 创建文件夹
+    image_dir = "original_images"
+    mask_dir = "mask_images"
+    os.makedirs(image_dir, exist_ok=True)
+    os.makedirs(mask_dir, exist_ok=True)
+
+    # 遍历数据集并保存图像
+    for idx, example in enumerate(dataset):
+        # 保存原图
+        image_path = os.path.join(image_dir, f"{idx:04d}.png")
+        example["Wang_Leehom_poster_image"].save(image_path)
+
+        # 保存 mask 图像
+        mask_path = os.path.join(mask_dir, f"{idx:04d}.png")
+        example["mask_image"].convert("RGB").save(mask_path)
+
+    print(f"✅ 原图已保存至 {image_dir}")
+    print(f"✅ Mask 图像已保存至 {mask_dir}")
+
+if __name__ == "__main__":
+    save_images_locally()
+
+import cv2
+import numpy as np
+from PIL import Image
+
+def blur_mask_opencv(
+    mask: Image.Image, 
+    blur_factor: int, 
+    edge_expand: int = 0  # 新增参数：边缘扩展像素值
+) -> Image.Image:
+    """
+    增强版遮罩处理：支持边缘扩展 + 高斯模糊羽化
+    参数:
+        mask: PIL格式的遮罩图像（支持 RGBA/RGB/L 模式）
+        blur_factor: 模糊强度（奇数）
+        edge_expand: 边缘扩展像素值（非负整数）
+    """
+    # 转换为OpenCV格式（NumPy数组）
+    mask_np = np.array(mask)
+    original_mode = mask.mode  # 保存原始色彩模式
+    
+    # 通道分离逻辑
+    if original_mode == 'RGBA':
+        alpha = mask_np[:, :, 3]
+        processed = alpha
+    elif original_mode == 'RGB':
+        processed = cv2.cvtColor(mask_np, cv2.COLOR_RGB2GRAY)
+    else:  # L模式（灰度）
+        processed = mask_np
+    
+    # 边缘扩展（核心新增功能）
+    if edge_expand > 0:
+        kernel = np.ones((3, 3), np.uint8)  # 3x3方形结构核
+        processed = cv2.dilate(
+            processed, 
+            kernel, 
+            iterations=edge_expand  # 扩展次数 = 扩展像素值
+        )
+    
+    # 高斯模糊羽化
+    kernel_size = blur_factor | 1  # 确保为奇数
+    blurred = cv2.GaussianBlur(processed, (kernel_size, kernel_size), 0)
+    
+    # 重建为PIL图像
+    if original_mode == 'RGBA':
+        result = mask_np.copy()
+        result[:, :, 3] = blurred  # 仅替换Alpha通道
+        return Image.fromarray(result)
+    else:
+        return Image.fromarray(blurred)
+
+
+#from PIL import Image 
+#mask = Image.open("image (51).jpg")
+
+# 使用示例
+#blurred_mask = blur_mask_opencv(mask, 33, 20)
+#blurred_mask.save("b_mask.png")
+#blurred_mask
+
+
+import os
+import numpy as np
+from PIL import Image
+import torch
+from diffusers import FluxFillPipeline
+
+# 初始化模型
+pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload()
+
+def calculate_optimal_dimensions(image):
+    original_width, original_height = image.size
+    MIN_ASPECT_RATIO = 9 / 16
+    MAX_ASPECT_RATIO = 16 / 9
+    FIXED_DIMENSION = 1024
+    original_aspect_ratio = original_width / original_height
+
+    if original_aspect_ratio > 1:
+        width = FIXED_DIMENSION
+        height = round(FIXED_DIMENSION / original_aspect_ratio)
+    else:
+        height = FIXED_DIMENSION
+        width = round(FIXED_DIMENSION * original_aspect_ratio)
+
+    width = (width // 8) * 8
+    height = (height // 8) * 8
+    calculated_aspect_ratio = width / height
+
+    if calculated_aspect_ratio > MAX_ASPECT_RATIO:
+        width = (height * MAX_ASPECT_RATIO // 8) * 8
+    elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
+        height = (width / MIN_ASPECT_RATIO // 8) * 8
+
+    width = max(width, 576) if width == FIXED_DIMENSION else width
+    height = max(height, 576) if height == FIXED_DIMENSION else height
+    return width, height
+
+def calculate_white_pixel_ratio(mask_image):
+    mask_array = np.array(mask_image.convert('L'))
+    white_pixels = np.sum(mask_array == 255)
+    total_pixels = mask_array.size
+    return white_pixels / total_pixels
+
+def inpaint(image_path, mask_path, output_path):
+    image = Image.open(image_path).convert("RGB")
+    mask = Image.open(mask_path).convert("L")
+    #mask = blur_mask_opencv(mask, 33, 20)
+    mask = blur_mask_opencv(mask, 33, 60)
+
+    white_ratio = calculate_white_pixel_ratio(mask)
+    if white_ratio >= 1 / 3:
+        image.save(output_path)
+        return
+
+    width, height = calculate_optimal_dimensions(image)
+    result = pipe(
+        prompt="",
+        height=height,
+        width=width,
+        image=image,
+        mask_image=mask,
+        num_inference_steps=40,
+        guidance_scale=28,
+    ).images[0]
+    result.save(output_path)
+
+def process_all_image_pairs():
+    image_dir = "original_images"
+    mask_dir = "mask_images"
+    output_dir = "inpaint_results"
+    os.makedirs(output_dir, exist_ok=True)
+
+    for filename in os.listdir(image_dir):
+        base_name = os.path.splitext(filename)[0]
+        image_path = os.path.join(image_dir, filename)
+        mask_path = os.path.join(mask_dir, filename)
+        output_path = os.path.join(output_dir, filename)
+
+        if not os.path.exists(mask_path):
+            continue
+
+        try:
+            inpaint(image_path, mask_path, output_path)
+        except Exception as e:
+            print(f"❌ 处理 {filename} 时出错: {e}")
+
+    print(f"✅ 修复结果已保存至 {output_dir}")
+
+if __name__ == "__main__":
+    process_all_image_pairs()
+
+import os
+from datasets import load_dataset, DatasetDict, Image as HfImage
+from PIL import Image
+
+def update_dataset_with_inpaint_results():
+    dataset_name = "svjack/InfiniteYou_PosterCraft_Wang_Leehom_Poster_FP8_WAV_text_mask"
+    input_dataset = load_dataset(dataset_name)["train"]
+
+    output_dir = "inpaint_results"
+    output_dataset_path = "InfiniteYou_PosterCraft_Wang_Leehom_Poster_FP8_WAV_text_mask_inpaint"
+    os.makedirs(output_dataset_path, exist_ok=True)
+
+    def add_inpaint_image(example, idx):
+        output_path = os.path.join(output_dir, f"{idx:04d}.png")
+        
+        if os.path.exists(output_path):
+            example["inpaint_image"] = output_path
+        else:
+            # 如果没有生成修复图像，则使用原始图像
+            example["inpaint_image"] = example["Wang_Leehom_poster_image"]
+        
+        return example
+
+    updated_dataset = input_dataset.map(
+        lambda ex, idx: add_inpaint_image(ex, idx),
+        with_indices=True,
+        batched=False,
+        num_proc=1
+    )
+
+    updated_dataset = updated_dataset.cast_column("inpaint_image", HfImage())
+    updated_dataset.save_to_disk(output_dataset_path)
+
+    print(f"✅ 更新后的数据集已保存至 {output_dataset_path}")
+
+if __name__ == "__main__":
+    update_dataset_with_inpaint_results()
+'''
+
 import torch
 import spaces
 import gradio as gr