草图数据生成方法整理

# download annotations (~275Mb)
wget https://dl.fbaipublicfiles.com/amateur_drawings/amateur_drawings_annotations.json

# download images (~50Gb)
wget https://dl.fbaipublicfiles.com/amateur_drawings/amateur_drawings.tar

[Category]
    - {obj}: 3D models
    - {mv-img}: Multi-view reference images rendered from 3D models
    - {tracing}: The tracings over the reference images
    - [Type]_png: The drawings rasterized from .json files
    - [Type]_json: The sequence data includes pressure, timeslot, drawing path, etc.

* [Category] indicates the category name, e.g. {Animal_Head}.
* [Type] indicates {original} type (user drawings) or registered types ({gloabl} -> sketch-level, {stroke} -> stroke-level and {reg} -> pixel-level) from the user drawings

import svgwrite
from svgpathtools import svg2paths2, wsvg
import numpy as np
import os
import cv2
import cairosvg
from PIL import Image
from svgpathtools import svg2paths2
from svgpathtools import Path as SVGPath

def extract_edge_distance_map(png_path, size=(512, 512), white_threshold=240):
    """
    提取边缘距离图，支持RGBA和RGB图片
    
    参数:
        png_path: 输入图片路径
        size: 目标尺寸
        white_threshold: 判定白色的阈值(0-255)
    """
    # Load image
    img = Image.open(png_path)
    img = img.resize(size)
    
    if img.mode == "RGBA":
        # 处理带透明通道的图片
        alpha = np.array(img)[:, :, 3]
        mask = (alpha > 0).astype(np.uint8) * 255
    else:
        # 处理不带透明通道的图片
        img_array = np.array(img.convert("RGB"))
        # 判断白色像素
        is_white = np.all(img_array >= white_threshold, axis=2)
        mask = (~is_white).astype(np.uint8) * 255
    
    # 保存并返回mask
    Image.fromarray(mask, mode='L').save("mask_map.png")
    
    # 计算距离变换
    distance_map = cv2.distanceTransform(
        mask, 
        distanceType=cv2.DIST_L2,  # 使用欧氏距离
        maskSize=cv2.DIST_MASK_PRECISE  # 使用精确计算
    )
    
    # 归一化距离图到0-255范围
    distance_map = cv2.normalize(
        distance_map, 
        None, 
        0, 
        255, 
        cv2.NORM_MINMAX, 
        dtype=cv2.CV_8U
    )
    # 保存距离图
    Image.fromarray(distance_map).save("distance_map.png")    
    return distance_map


def sample_path_points(path: SVGPath, num_points=100):
    """Sample `num_points` along an SVG path."""
    return [path.point(t) for t in np.linspace(0, 1, num_points)]

def normalize_svg_point(point, viewbox, target_size=(512, 512)):
    """将 SVG 坐标标准化为 512x512 图像坐标"""
    min_x, min_y, width, height = viewbox
    x = (point.real - min_x) / width * target_size[0]
    y = (point.imag - min_y) / height * target_size[1]
    return int(x), int(y)

def process_distances(valid_distances, method='minmax'):
    """
    处理距离数据，使其分布更均匀
    
    参数:
        valid_distances: 有效距离列表
        method: 处理方法 ('log', 'robust', 'minmax', 'zscore')
    """
    distances = np.array(valid_distances)
    
    if method == 'log':
        # 对数变换
        processed = np.log1p(distances)
    
    elif method == 'robust':
        # 稳健缩放（基于分位数）
        q1, q3 = np.percentile(distances, [25, 75])
        iqr = q3 - q1
        lower_bound = q1 - 1.5 * iqr
        upper_bound = q3 + 1.5 * iqr
        
        # 过滤异常值并缩放
        mask = (distances >= lower_bound) & (distances <= upper_bound)
        processed = distances[mask]
        
        # 重新缩放到0-1范围
        processed = (processed - processed.min()) / (processed.max() - processed.min())
        
    elif method == 'minmax':
        # 最小最大缩放
        processed = (distances - distances.min()) / (distances.max() - distances.min())
        
    elif method == 'zscore':
        # Z-score标准化
        mean = np.mean(distances)
        std = np.std(distances)
        processed = (distances - mean) / std
    
    return processed


def filter_by_edge_and_length(svg_path, png_path, k=0.5, s=0.5, 
                              mode='and', output_svg_path="filtered_final.svg", 
                              canvas_size=(512, 512), soft_edge=True):
    from svgpathtools import svg2paths2, wsvg

    # 读取路径与 SVG 属性
    paths, attributes, svg_attr = svg2paths2(svg_path)
    viewbox = [float(x) for x in svg_attr['viewBox'].split()]  # [min_x, min_y, width, height]

    # 距离图
    distance_map = extract_edge_distance_map(png_path, size=canvas_size)

    edge_filtered_set = set()
    length_filtered_set = set()
    all_path_distances = []
    all_path_lengths = []

    for i, path in enumerate(paths):
        # ---- 计算平均距离 ----
        points = sample_path_points(path, num_points=200)
        distances = []
        for p in points:
            x, y = normalize_svg_point(p, viewbox, canvas_size)
            if 0 <= x < canvas_size[0] and 0 <= y < canvas_size[1]:
                distances.append(distance_map[y, x])
        avg_distance = np.mean(distances) if distances else None
        all_path_distances.append(avg_distance)

        # ---- 计算路径长度 ----
        length = path.length() * 10
        all_path_lengths.append(length)

    # 计算距离阈值
    all_path_distances = process_distances(all_path_distances, method='log')
    valid_distances = [d for d in all_path_distances if d is not None]
    a_d, b_d = min(valid_distances), max(valid_distances)
    distance_threshold = a_d + k * (b_d - a_d)
    # print(f"max={b_d}, min={a_d}, thres={distance_threshold}", end='')
    # print(len(valid_distances), valid_distances)

    # 计算长度阈值
    all_path_lengths = process_distances(all_path_lengths)
    a_l, b_l = max(all_path_lengths), min(all_path_lengths)
    length_threshold = b_l + s * (a_l - b_l)
    
    # lengths = process_distances(np.array([l for l in all_path_lengths if l is not None]))
    # mean_length = np.mean(lengths)  # 平均长度
    # std_length = np.std(lengths)    # 标准差
    # length_threshold = mean_length + s * std_length
    # print(len(all_path_lengths), all_path_lengths)

    # 满足距离条件的路径索引
    nums = 0
    for idx, (d, l) in enumerate(zip(all_path_distances, all_path_lengths)):
        if l > b_l + 0.05 * (a_l - b_l):
            nums += 1
        if l >= length_threshold:
                length_filtered_set.add(idx)
        if d is not None and d <= distance_threshold:
            if soft_edge and l <= b_l + 0.1 * (a_l - b_l):
                continue
            edge_filtered_set.add(idx)
    print(f"nums = {nums}, ", end='')


    # 合并结果
    if mode == 'and':
        final_indices = edge_filtered_set & length_filtered_set
    elif mode == 'or':
        final_indices = edge_filtered_set | length_filtered_set
    else:
        raise ValueError("mode 应为 'and' 或 'or'")

    # 生成新的路径列表
    filtered_paths = [paths[i] for i in final_indices]
    filtered_attributes = [attributes[i] for i in final_indices]

    # 保存新 SVG
    wsvg(filtered_paths, attributes=filtered_attributes, svg_attributes=svg_attr, filename=output_svg_path)
    print(f"共保留 {len(filtered_paths)} 条路径，保存至：{output_svg_path}")


def create_image_grid_enhanced(svg_dir, output_path, ks, ss,
                             single_size=(256, 256), padding=10, 
                             background_color='white'):
    """
    网格图像创建函数
    
    功能：
    - 图像间距
    - 背景颜色选择
    - 可选的标签
    """
    grid_size = (len(ks), len(ss))
    rows, cols = grid_size
    width, height = single_size
    
    # 计算含padding的总尺寸
    total_width = width * cols + padding * (cols + 1)
    total_height = height * rows + padding * (rows + 1)
    
    # 创建空白画布
    grid_image = Image.new('RGB', 
                          (total_width, total_height), 
                          color=background_color)
    
    print(f"final ks = {ks}   , ss = {ss}")
    for i, k in enumerate(ks):
        for j, s in enumerate(ss):
            svg_name = f"k_{k}_s_{s}.svg"
            svg_path = os.path.join(svg_dir, svg_name)
            temp_png = f"temp_{i}_{j}.png"
            
            try:
                # 转换SVG到PNG
                cairosvg.svg2png(
                    url=svg_path,
                    write_to=temp_png,
                    output_width=width,
                    output_height=height,
                    # background_color="white" 
                )
                
                img = Image.open(temp_png)
                
                # 计算带padding的位置
                x = j * (width + padding) + padding
                y = i * (height + padding) + padding
                
                # 粘贴到网格中
                grid_image.paste(img, (x, y))
                
                os.remove(temp_png)
                
            except Exception as e:
                print(f"处理 {svg_name} 时出错: {e}")
    
    grid_image.save(output_path)
    print(f"网格图像已保存到 {output_path}")


if __name__ == '__main__':
    svg_path = f'xxx.svg'
    png_path = f'mask.png'  # 用于提取边缘图
    output_path = "./assets/diff/svg"

    ks = [0.9, 0.7, 0.5, 0.3]   # 距离
    ss = [0.1, 0.2, 0.3, 0.4]   # 长度
    for k in ks:
        for s in ss:
            out_path = os.path.join(output_path, f"k_{k}_s_{s}.svg")
            print(f"k = {k}, s = {s}, out_path = {out_path}...")
            filter_by_edge_and_length(svg_path, png_path, k=k, s=s, mode='or', output_svg_path=out_path)

    print("generate the final image...")
    create_image_grid_enhanced(output_path, f".diff/tot.png", ks, ss)