Upload 7 files

27d0379 verified 15 days ago

12.2 kB

import os
import glob
import subprocess
import pandas as pd
from datasets import Dataset, Features, Image as HFImage, Value
from huggingface_hub import HfApi, DatasetCard
import tempfile
from PIL import Image

# ==========================================
# 配置区
# ==========================================
REPO_ID = "waveshaker/dataset_collections" 
REPO_TYPE = "dataset"
HF_TOKEN = ""  # 请填入你的 Token
# ==========================================

def clean_txt_newlines(current_dir):
    """1. 递归清理所有 txt 文件中的换行符"""
    print("\n" + "="*50)
    print("🧹 [任务 1/5] 正在清理所有 TXT 文件中的换行符...")
    print("="*50)
    
    txt_files = glob.glob(os.path.join(current_dir, "**", "*.txt"), recursive=True)
    if not txt_files:
        print("  ✅ 没有找到需要处理的 TXT 文件。")
        return

    cleaned_count = 0
    for txt_path in txt_files:
        # 跳过根目录下的 txt（可能是说明文档）
        if os.path.dirname(txt_path) == current_dir:
            continue
            
        try:
            with open(txt_path, 'r', encoding='utf-8') as f:
                content = ''.join(line.strip() for line in f)
            
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(content)
            cleaned_count += 1
        except Exception as e:
            print(f"  ⚠️ 处理 {os.path.basename(txt_path)} 失败: {e}")
            
    print(f"  🎉 TXT 清理完成！共处理了 {cleaned_count} 个文件。")

def remove_mask_metadata(current_dir):
    """2. 递归剥离 masklabel.png 的元数据，保护隐私"""
    print("\n" + "="*50)
    print("🛡️ [任务 2/5] 正在剥离 Mask 蒙版图 (masklabel.png) 的元数据...")
    print("="*50)
    
    png_files = glob.glob(os.path.join(current_dir, "**", "*.png"), recursive=True)
    mask_files = [f for f in png_files if "masklabel" in os.path.basename(f).lower()]
    
    if not mask_files:
        print("  ✅ 没有找到 masklabel.png 文件，跳过此步骤。")
        return

    success_count = 0
    for file_path in mask_files:
        try:
            with Image.open(file_path) as img:
                # 读取像素数据，彻底断开原元数据
                data = list(img.getdata())
                clean_img = Image.new(img.mode, img.size)
                clean_img.putdata(data)
                # 覆盖保存
                clean_img.save(file_path, "PNG")
            success_count += 1
            print(f"  🔒 已清理元数据: {os.path.basename(file_path)}")
        except Exception as e:
            print(f"  ⚠️ 处理失败 [{os.path.basename(file_path)}]: {e}")

    print(f"  🎉 元数据清理完成！共处理了 {success_count} 张 Mask 图。")

def convert_png_to_jpg(current_dir):
    """3. 批量转换常规 png 为 jpg，并删除原 png（绝对跳过 masklabel）"""
    print("\n" + "="*50)
    print("🎬 [任务 3/5] 正在批量转换常规 PNG 为 JPG...")
    print("="*50)
    
    png_files = glob.glob(os.path.join(current_dir, "**", "*.png"), recursive=True)
    if not png_files:
        print("  ✅ 没有找到需要转换的 PNG 文件。")
        return

    converted_count = 0
    for png_path in png_files:
        # 🚨 核心过滤：跳过 masklabel 文件，保证蒙版依然是 png
        if "masklabel" in png_path.lower():
            continue
            
        jpg_path = os.path.splitext(png_path)[0] + ".jpg"
        
        try:
            subprocess.run(
                ["ffmpeg", "-y", "-i", png_path, "-q:v", "2", jpg_path],
                stdout=subprocess.DEVNULL, 
                stderr=subprocess.DEVNULL,
                check=True
            )
            os.remove(png_path)
            converted_count += 1
            print(f"  🔄 转换: {os.path.basename(png_path)} -> {os.path.basename(jpg_path)}")
        except subprocess.CalledProcessError:
            print(f"  ⚠️ 转换失败: {png_path}")
        except FileNotFoundError:
            print("  ❌ 错误: 找不到 ffmpeg！请确保已安装并配置了环境变量。")
            return

    print(f"  🎉 转换完成！共转换 {converted_count} 张图片。")

def generate_parquet_and_update_yaml(api, current_dir, region_prefix):
    """4. 本地打包生成 parquet，拉取云端 YAML，追加并重排"""
    print("\n" + "="*50)
    print("📦 [任务 4/5] 正在本地打包 Parquet 并更新 YAML 配置...")
    print("="*50)
    
    local_folder_name = os.path.basename(current_dir).replace(" ", "_").replace("-", "_")
    
    print(f"  ☁️ 正在读取云端 {REPO_ID} 的现有配置...")
    try:
        remote_readme_path = api.hf_hub_download(repo_id=REPO_ID, filename="README.md", repo_type=REPO_TYPE)
        with open(remote_readme_path, "r", encoding="utf-8") as f:
            content = f.read()
        card = DatasetCard(content)
        print("  ✅ 成功读取云端 README。")
    except Exception as e:
        print(f"  ⚠️ 找不到云端 README，将新建: {e}")
        card = DatasetCard("")

    yaml_data = card.data if card.data else {}
    if "configs" not in yaml_data: yaml_data["configs"] = []
    
    yaml_data["dataset_info"] = {
        "features": [
            {"name": "image", "dtype": "image"},
            {"name": "file_name", "dtype": "string"},
            {"name": "text", "dtype": "string"}
        ]
    }

    existing_config_names = [c.get("config_name") for c in yaml_data["configs"]]
    new_configs_added = False
    valid_folders = []

    for sub_folder_name in os.listdir(current_dir):
        sub_folder_path = os.path.join(current_dir, sub_folder_name)
        
        if os.path.isdir(sub_folder_path) and not sub_folder_name.startswith('.'):
            jpg_files = glob.glob(os.path.join(sub_folder_path, "*.jpg")) + glob.glob(os.path.join(sub_folder_path, "*.jpeg"))
            if not jpg_files:
                continue
                
            print(f"  📦 正在打包: {sub_folder_name}")
            data_rows = []
            
            for jpg_path in jpg_files:
                file_name = os.path.basename(jpg_path)
                base_name = os.path.splitext(file_name)[0]
                txt_path = os.path.join(sub_folder_path, f"{base_name}.txt")
                
                text_content = ""
                if os.path.exists(txt_path):
                    with open(txt_path, "r", encoding="utf-8") as f:
                        text_content = f.read().strip()

                try:
                    with open(jpg_path, "rb") as f:
                        image_bytes = {"bytes": f.read(), "path": None}
                except Exception as e:
                    print(f"    ⚠️ 跳过损坏图片 {file_name}: {e}")
                    continue

                data_rows.append({
                    "image": image_bytes,
                    "file_name": file_name,
                    "text": text_content
                })

            if data_rows:
                df = pd.DataFrame(data_rows)
                features = Features({
                    "image": HFImage(),  # 👈 这里加上了 HF
                    "file_name": Value("string"),
                    "text": Value("string")
                })
                ds = Dataset.from_pandas(df, features=features)
                ds.to_parquet(os.path.join(sub_folder_path, "train.parquet"))
                valid_folders.append(sub_folder_name)
                print(f"    ✅ 生成 train.parquet ({len(data_rows)} 条数据)")

                config_name = f"{local_folder_name}_{sub_folder_name}".lower()
                path_str = f"{region_prefix}/{local_folder_name}/{sub_folder_name}/train.parquet"
                
                if config_name not in existing_config_names:
                    yaml_data["configs"].append({
                        "config_name": config_name,
                        "data_files": [{"split": "train", "path": path_str}]
                    })
                    new_configs_added = True
                    print(f"    ➕ 已添加配置: {config_name}")

    def get_priority(config):
        data_files = config.get("data_files", [])
        path = data_files[0].get("path", "") if data_files else ""
        # 直接使用原始大小写，确保与云端路径完全一致
        priority_dict = {"CN": 1, "KR": 2, "JP": 3, "Fictional": 4, "WE": 5}
        path_prefix = path.split('/')[0] if path else ""
        return priority_dict.get(path_prefix, 99)

    yaml_data["configs"] = sorted(yaml_data["configs"], key=lambda c: (get_priority(c), c.get("config_name", "")))
    print("  🔄 已根据 [CN > KR > JP > Fictional > WE] 的顺序重排 YAML。")

    if new_configs_added:
        card.data = yaml_data
        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', suffix='.md', delete=False) as temp_readme:
            temp_readme.write(str(card))
            return temp_readme.name, valid_folders
    return None, valid_folders

def upload_files_and_readme(api, current_dir, region_prefix, temp_readme_path, valid_folders):
    """5. 上传文件及 README"""
    print("\n" + "="*50)
    print("🚀 [任务 5/5] 正在上传文件到 Hugging Face...")
    print("="*50)
    
    local_folder_name = os.path.basename(current_dir).replace(" ", "_").replace("-", "_")
    
    for sub_folder_name in valid_folders:
        sub_folder_path = os.path.join(current_dir, sub_folder_name)
        path_in_repo = f"{region_prefix}/{local_folder_name}/{sub_folder_name}"
        
        print(f"  📦 开始上传: {sub_folder_name} -> {path_in_repo}")
        try:
            api.upload_folder(
                folder_path=sub_folder_path,
                repo_id=REPO_ID,
                repo_type=REPO_TYPE,
                path_in_repo=path_in_repo,
                commit_message=f"Update files for {path_in_repo}",
                # 上传过滤：保留 masklabel.png，jpg，txt，parquet
                allow_patterns=["*.jpg", "*.jpeg", "*.txt", "*.png", "*.parquet"]
            )
            print(f"  🎉 {sub_folder_name} 上传完成！")
        except Exception as e:
            print(f"  ❌ {sub_folder_name} 上传失败: {e}")

    if temp_readme_path:
        print("\n  📝 正在推送最新的 README.md...")
        try:
            api.upload_file(
                path_or_fileobj=temp_readme_path,
                path_in_repo="README.md",
                repo_id=REPO_ID,
                repo_type=REPO_TYPE,
                commit_message="Auto-update README YAML configs"
            )
            print("  🎉 README.md 已成功更新到云端！")
        except Exception as e:
            print(f"  ❌ README.md 上传失败: {e}")
        finally:
            os.remove(temp_readme_path)

def main():
    api = HfApi(token=HF_TOKEN)
    current_dir = os.getcwd() 
    
    print("\n" + "="*50)
    print("🌟 欢迎使用 5合1 数据集自动化流水线 🌟")
    print("="*50)
    
    print("请选择该数据集的地区/分类前缀：")
    print("1. CN | 2. KR | 3. JP | 4. Fictional | 5. WE | 6. 自定义")
    choice = input("\n👉 请输入数字或直接输入前缀 (默认回车为 KR): ").strip().upper()
    if not choice: choice = "2"
    mapping = {"1": "CN", "2": "KR", "3": "JP", "4": "Fictional", "5": "WE"}
    region_prefix = mapping.get(choice, choice).replace(" ", "_").replace("-", "_")
    print(f"\n✅ 已选择前缀: {region_prefix}")

    # 任务 1：清理 txt 换行符
    clean_txt_newlines(current_dir)

    # 任务 2：清理 Mask 元数据
    remove_mask_metadata(current_dir)

    # 任务 3：转换 PNG -> JPG
    convert_png_to_jpg(current_dir)

    # 任务 4：打包 Parquet 并更新 YAML
    temp_readme_path, valid_folders = generate_parquet_and_update_yaml(api, current_dir, region_prefix)

    # 任务 5：上传
    upload_files_and_readme(api, current_dir, region_prefix, temp_readme_path, valid_folders)

    print("\n" + "="*50)
    print("🎉 5合1 全流程处理完毕！")
    print("="*50)

if __name__ == "__main__":
    main()

	import os
	import glob
	import subprocess
	import pandas as pd
	from datasets import Dataset, Features, Image as HFImage, Value
	from huggingface_hub import HfApi, DatasetCard
	import tempfile
	from PIL import Image

	# ==========================================
	# 配置区
	# ==========================================
	REPO_ID = "waveshaker/dataset_collections"
	REPO_TYPE = "dataset"
	HF_TOKEN = "" # 请填入你的 Token
	# ==========================================

	def clean_txt_newlines(current_dir):
	"""1. 递归清理所有 txt 文件中的换行符"""
	print("\n" + "="*50)
	print("🧹 [任务 1/5] 正在清理所有 TXT 文件中的换行符...")
	print("="*50)

	txt_files = glob.glob(os.path.join(current_dir, "*", ".txt"), recursive=True)
	if not txt_files:
	print(" ✅ 没有找到需要处理的 TXT 文件。")
	return

	cleaned_count = 0
	for txt_path in txt_files:
	# 跳过根目录下的 txt（可能是说明文档）
	if os.path.dirname(txt_path) == current_dir:
	continue

	try:
	with open(txt_path, 'r', encoding='utf-8') as f:
	content = ''.join(line.strip() for line in f)

	with open(txt_path, 'w', encoding='utf-8') as f:
	f.write(content)
	cleaned_count += 1
	except Exception as e:
	print(f" ⚠️ 处理 {os.path.basename(txt_path)} 失败: {e}")

	print(f" 🎉 TXT 清理完成！共处理了 {cleaned_count} 个文件。")

	def remove_mask_metadata(current_dir):
	"""2. 递归剥离 masklabel.png 的元数据，保护隐私"""
	print("\n" + "="*50)
	print("🛡️ [任务 2/5] 正在剥离 Mask 蒙版图 (masklabel.png) 的元数据...")
	print("="*50)

	png_files = glob.glob(os.path.join(current_dir, "*", ".png"), recursive=True)
	mask_files = [f for f in png_files if "masklabel" in os.path.basename(f).lower()]

	if not mask_files:
	print(" ✅ 没有找到 masklabel.png 文件，跳过此步骤。")
	return

	success_count = 0
	for file_path in mask_files:
	try:
	with Image.open(file_path) as img:
	# 读取像素数据，彻底断开原元数据
	data = list(img.getdata())
	clean_img = Image.new(img.mode, img.size)
	clean_img.putdata(data)
	# 覆盖保存
	clean_img.save(file_path, "PNG")
	success_count += 1
	print(f" 🔒 已清理元数据: {os.path.basename(file_path)}")
	except Exception as e:
	print(f" ⚠️ 处理失败 [{os.path.basename(file_path)}]: {e}")

	print(f" 🎉 元数据清理完成！共处理了 {success_count} 张 Mask 图。")

	def convert_png_to_jpg(current_dir):
	"""3. 批量转换常规 png 为 jpg，并删除原 png（绝对跳过 masklabel）"""
	print("\n" + "="*50)
	print("🎬 [任务 3/5] 正在批量转换常规 PNG 为 JPG...")
	print("="*50)

	png_files = glob.glob(os.path.join(current_dir, "*", ".png"), recursive=True)
	if not png_files:
	print(" ✅ 没有找到需要转换的 PNG 文件。")
	return

	converted_count = 0
	for png_path in png_files:
	# 🚨 核心过滤：跳过 masklabel 文件，保证蒙版依然是 png
	if "masklabel" in png_path.lower():
	continue

	jpg_path = os.path.splitext(png_path)[0] + ".jpg"

	try:
	subprocess.run(
	["ffmpeg", "-y", "-i", png_path, "-q:v", "2", jpg_path],
	stdout=subprocess.DEVNULL,
	stderr=subprocess.DEVNULL,
	check=True
	)
	os.remove(png_path)
	converted_count += 1
	print(f" 🔄 转换: {os.path.basename(png_path)} -> {os.path.basename(jpg_path)}")
	except subprocess.CalledProcessError:
	print(f" ⚠️ 转换失败: {png_path}")
	except FileNotFoundError:
	print(" ❌ 错误: 找不到 ffmpeg！请确保已安装并配置了环境变量。")
	return

	print(f" 🎉 转换完成！共转换 {converted_count} 张图片。")

	def generate_parquet_and_update_yaml(api, current_dir, region_prefix):
	"""4. 本地打包生成 parquet，拉取云端 YAML，追加并重排"""
	print("\n" + "="*50)
	print("📦 [任务 4/5] 正在本地打包 Parquet 并更新 YAML 配置...")
	print("="*50)

	local_folder_name = os.path.basename(current_dir).replace(" ", "_").replace("-", "_")

	print(f" ☁️ 正在读取云端 {REPO_ID} 的现有配置...")
	try:
	remote_readme_path = api.hf_hub_download(repo_id=REPO_ID, filename="README.md", repo_type=REPO_TYPE)
	with open(remote_readme_path, "r", encoding="utf-8") as f:
	content = f.read()
	card = DatasetCard(content)
	print(" ✅ 成功读取云端 README。")
	except Exception as e:
	print(f" ⚠️ 找不到云端 README，将新建: {e}")
	card = DatasetCard("")

	yaml_data = card.data if card.data else {}
	if "configs" not in yaml_data: yaml_data["configs"] = []

	yaml_data["dataset_info"] = {
	"features": [
	{"name": "image", "dtype": "image"},
	{"name": "file_name", "dtype": "string"},
	{"name": "text", "dtype": "string"}
	]
	}

	existing_config_names = [c.get("config_name") for c in yaml_data["configs"]]
	new_configs_added = False
	valid_folders = []

	for sub_folder_name in os.listdir(current_dir):
	sub_folder_path = os.path.join(current_dir, sub_folder_name)

	if os.path.isdir(sub_folder_path) and not sub_folder_name.startswith('.'):
	jpg_files = glob.glob(os.path.join(sub_folder_path, ".jpg")) + glob.glob(os.path.join(sub_folder_path, ".jpeg"))
	if not jpg_files:
	continue

	print(f" 📦 正在打包: {sub_folder_name}")
	data_rows = []

	for jpg_path in jpg_files:
	file_name = os.path.basename(jpg_path)
	base_name = os.path.splitext(file_name)[0]
	txt_path = os.path.join(sub_folder_path, f"{base_name}.txt")

	text_content = ""
	if os.path.exists(txt_path):
	with open(txt_path, "r", encoding="utf-8") as f:
	text_content = f.read().strip()

	try:
	with open(jpg_path, "rb") as f:
	image_bytes = {"bytes": f.read(), "path": None}
	except Exception as e:
	print(f" ⚠️ 跳过损坏图片 {file_name}: {e}")
	continue

	data_rows.append({
	"image": image_bytes,
	"file_name": file_name,
	"text": text_content
	})

	if data_rows:
	df = pd.DataFrame(data_rows)
	features = Features({
	"image": HFImage(), # 👈 这里加上了 HF
	"file_name": Value("string"),
	"text": Value("string")
	})
	ds = Dataset.from_pandas(df, features=features)
	ds.to_parquet(os.path.join(sub_folder_path, "train.parquet"))
	valid_folders.append(sub_folder_name)
	print(f" ✅ 生成 train.parquet ({len(data_rows)} 条数据)")

	config_name = f"{local_folder_name}_{sub_folder_name}".lower()
	path_str = f"{region_prefix}/{local_folder_name}/{sub_folder_name}/train.parquet"

	if config_name not in existing_config_names:
	yaml_data["configs"].append({
	"config_name": config_name,
	"data_files": [{"split": "train", "path": path_str}]
	})
	new_configs_added = True
	print(f" ➕ 已添加配置: {config_name}")

	def get_priority(config):
	data_files = config.get("data_files", [])
	path = data_files[0].get("path", "") if data_files else ""
	# 直接使用原始大小写，确保与云端路径完全一致
	priority_dict = {"CN": 1, "KR": 2, "JP": 3, "Fictional": 4, "WE": 5}
	path_prefix = path.split('/')[0] if path else ""
	return priority_dict.get(path_prefix, 99)

	yaml_data["configs"] = sorted(yaml_data["configs"], key=lambda c: (get_priority(c), c.get("config_name", "")))
	print(" 🔄 已根据 [CN > KR > JP > Fictional > WE] 的顺序重排 YAML。")

	if new_configs_added:
	card.data = yaml_data
	with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', suffix='.md', delete=False) as temp_readme:
	temp_readme.write(str(card))
	return temp_readme.name, valid_folders
	return None, valid_folders

	def upload_files_and_readme(api, current_dir, region_prefix, temp_readme_path, valid_folders):
	"""5. 上传文件及 README"""
	print("\n" + "="*50)
	print("🚀 [任务 5/5] 正在上传文件到 Hugging Face...")
	print("="*50)

	local_folder_name = os.path.basename(current_dir).replace(" ", "_").replace("-", "_")

	for sub_folder_name in valid_folders:
	sub_folder_path = os.path.join(current_dir, sub_folder_name)
	path_in_repo = f"{region_prefix}/{local_folder_name}/{sub_folder_name}"

	print(f" 📦 开始上传: {sub_folder_name} -> {path_in_repo}")
	try:
	api.upload_folder(
	folder_path=sub_folder_path,
	repo_id=REPO_ID,
	repo_type=REPO_TYPE,
	path_in_repo=path_in_repo,
	commit_message=f"Update files for {path_in_repo}",
	# 上传过滤：保留 masklabel.png，jpg，txt，parquet
	allow_patterns=[".jpg", ".jpeg", ".txt", ".png", "*.parquet"]
	)
	print(f" 🎉 {sub_folder_name} 上传完成！")
	except Exception as e:
	print(f" ❌ {sub_folder_name} 上传失败: {e}")

	if temp_readme_path:
	print("\n 📝 正在推送最新的 README.md...")
	try:
	api.upload_file(
	path_or_fileobj=temp_readme_path,
	path_in_repo="README.md",
	repo_id=REPO_ID,
	repo_type=REPO_TYPE,
	commit_message="Auto-update README YAML configs"
	)
	print(" 🎉 README.md 已成功更新到云端！")
	except Exception as e:
	print(f" ❌ README.md 上传失败: {e}")
	finally:
	os.remove(temp_readme_path)

	def main():
	api = HfApi(token=HF_TOKEN)
	current_dir = os.getcwd()

	print("\n" + "="*50)
	print("🌟 欢迎使用 5合1 数据集自动化流水线 🌟")
	print("="*50)

	print("请选择该数据集的地区/分类前缀：")
	print("1. CN \| 2. KR \| 3. JP \| 4. Fictional \| 5. WE \| 6. 自定义")
	choice = input("\n👉 请输入数字或直接输入前缀 (默认回车为 KR): ").strip().upper()
	if not choice: choice = "2"
	mapping = {"1": "CN", "2": "KR", "3": "JP", "4": "Fictional", "5": "WE"}
	region_prefix = mapping.get(choice, choice).replace(" ", "_").replace("-", "_")
	print(f"\n✅ 已选择前缀: {region_prefix}")

	# 任务 1：清理 txt 换行符
	clean_txt_newlines(current_dir)

	# 任务 2：清理 Mask 元数据
	remove_mask_metadata(current_dir)

	# 任务 3：转换 PNG -> JPG
	convert_png_to_jpg(current_dir)

	# 任务 4：打包 Parquet 并更新 YAML
	temp_readme_path, valid_folders = generate_parquet_and_update_yaml(api, current_dir, region_prefix)

	# 任务 5：上传
	upload_files_and_readme(api, current_dir, region_prefix, temp_readme_path, valid_folders)

	print("\n" + "="*50)
	print("🎉 5合1 全流程处理完毕！")
	print("="*50)

	if __name__ == "__main__":
	main()