Spaces:
Runtime error
Runtime error
''' | |
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 | |
from diffusers import FluxFillPipeline | |
pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda") | |
# reference https://huggingface.co/spaces/black-forest-labs/FLUX.1-Fill-dev/blob/main/app.py | |
def calculate_optimal_dimensions(image): | |
# Extract the original dimensions | |
original_width, original_height = image.size | |
# Set constants | |
MIN_ASPECT_RATIO = 9 / 16 | |
MAX_ASPECT_RATIO = 16 / 9 | |
FIXED_DIMENSION = 1024 | |
# Calculate the aspect ratio of the original image | |
original_aspect_ratio = original_width / original_height | |
# Determine which dimension to fix | |
if original_aspect_ratio > 1: # Wider than tall | |
width = FIXED_DIMENSION | |
height = round(FIXED_DIMENSION / original_aspect_ratio) | |
else: # Taller than wide | |
height = FIXED_DIMENSION | |
width = round(FIXED_DIMENSION * original_aspect_ratio) | |
# Ensure dimensions are multiples of 8 | |
width = (width // 8) * 8 | |
height = (height // 8) * 8 | |
# Enforce aspect ratio limits | |
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 | |
# Ensure width and height remain above the minimum dimensions | |
width = max(width, 576) if width == FIXED_DIMENSION else width | |
height = max(height, 576) if height == FIXED_DIMENSION else height | |
return width, height | |
def inpaint( | |
image, | |
mask, | |
prompt="", | |
num_inference_steps=28, | |
guidance_scale=50, | |
): | |
image = image.convert("RGB") | |
mask = mask.convert("L") | |
width, height = calculate_optimal_dimensions(image) | |
result = pipe( | |
prompt=prompt, | |
height= height, | |
width= width, | |
image= image, | |
mask_image=mask, | |
num_inference_steps=num_inference_steps, | |
guidance_scale=guidance_scale, | |
).images[0] | |
result = result.convert("RGBA") | |
return result | |
demo = gr.Interface( | |
fn=inpaint, | |
inputs=[ | |
gr.Image(label="image", type="pil"), | |
gr.Image(label="mask", type="pil"), | |
gr.Text(label="prompt"), | |
gr.Number(value=40, label="num_inference_steps"), | |
gr.Number(value=28, label="guidance_scale"), | |
], | |
outputs=["image"], | |
api_name="inpaint", | |
examples=[["./assets/rocket.png", "./assets/Inpainting mask.png"]], | |
cache_examples=False, | |
description="it is recommended that you use https://github.com/la-voliere/react-mask-editor when creating an image mask in JS and then inverse it before sending it to this space", | |
) | |
#demo.launch(debug=True,show_error=True) | |
demo.launch(debug=True,show_error=True, share = True) |