Update app.py

app.py

@@ -1,3 +1,108 @@
+'''
+https://huggingface.co/spaces/merve/OWLSAM
+
+text,letter,watermark
+
+vim run_text_mask.py
+
+from gradio_client import Client, handle_file
+from datasets import load_dataset, Image as HfImage
+from PIL import ImageOps, Image
+import numpy as np
+import os
+from tqdm import tqdm
+
+# 初始化客户端
+client = Client("http://localhost:7860")
+
+# 加载数据集
+dataset_name = "svjack/InfiniteYou_PosterCraft_Wang_Leehom_Poster_FP8_WAV"
+dataset = load_dataset(dataset_name)
+
+# 创建保存 mask 的文件夹
+os.makedirs("mask_images", exist_ok=True)
+
+#### 832, 1216
+#### (864, 1152)
+def combine_non_white_regions(annotations):
+    canvas = None
+    for i, annotation in enumerate(annotations):
+        img = Image.open(annotation["image"]).convert("RGBA")
+        img_array = np.array(img)
+        if canvas is None:
+            height, width = img_array.shape[:2]
+            canvas = np.zeros((height, width, 4), dtype=np.uint8)
+        rgb = img_array[..., :3]
+        non_white_mask = np.any(rgb < 240, axis=-1, keepdims=True)
+        alpha_layer = np.where(non_white_mask, img_array[..., 3:], 0)
+        processed_img = np.concatenate([rgb, alpha_layer], axis=-1)
+        canvas = np.where(processed_img[..., 3:] > 0, processed_img, canvas)
+    if canvas is None:
+        height = 1152
+        width = 864
+        result_array = np.zeros((height, width, 4), dtype=np.uint8)
+        result_array[..., :3] = 255
+        result_array[..., 3] = 255
+        return Image.fromarray(result_array.astype(np.uint8))
+        
+    result_array = np.zeros((height, width, 4), dtype=np.uint8)
+    result_array[..., :3] = 255
+    result_array[..., 3] = 255
+    result_array = np.where(canvas[..., 3:] > 0, canvas, result_array)
+
+    non_white_mask = np.any(result_array[..., :3] < 255, axis=-1)
+    result_array[non_white_mask] = [0, 0, 0, 255]
+
+    return Image.fromarray(result_array.astype(np.uint8))
+
+def generate_mask(image, idx):
+    try:
+        # 保存原始图片为临时文件
+        temp_input_path = f"mask_images/temp_{idx:04d}.jpg"
+        image.save(temp_input_path)
+
+        # 调用 Gradio API
+        result = client.predict(
+            image=handle_file(temp_input_path),
+            texts="text,letter,watermark",
+            threshold=0.05,
+            sam_threshold=0.88,
+            api_name="/predict"
+        )
+
+        # 生成 mask 图像
+        mask_image = combine_non_white_regions(result["annotations"])
+        mask_image = ImageOps.invert(mask_image.convert("RGB"))
+
+        # 保存 mask 图像
+        output_mask_path = f"mask_images/mask_{idx:04d}.jpg"
+        mask_image.save(output_mask_path)
+
+        return {"mask_image": output_mask_path}
+
+    except Exception as e:
+        print(f"生成 mask 时出错 (index={idx}): {e}")
+        return {"mask_image": None}
+
+# 使用 map 处理整个数据集
+updated_dataset = dataset["train"].map(
+    lambda example, idx: generate_mask(example["Wang_Leehom_poster_image"], idx),
+    with_indices=True,
+    num_proc=1,
+    batched=False
+)
+
+# 转换列类型为 Image
+updated_dataset = updated_dataset.cast_column("mask_image", HfImage())
+
+# 保存更新后的数据集
+output_path = "Wang_Leehom_PosterCraft_with_Mask"
+updated_dataset.save_to_disk(output_path)
+
+print(f"✅ 已生成包含 mask 的数据集并保存至: {output_path}")
+'''
+
+
 from transformers import pipeline, SamModel, SamProcessor
 import torch
 import numpy as np