generate_dataset.py

import os
import asyncio
from pathlib import Path
from PIL import Image
import random
from minepi import Skin
from minepi.skin_render import Render

# 创建输出文件夹
RENDER_OUTPUT_DIR = "rendered_skins"
TEXTURE_OUTPUT_DIR = "scaled_textures"
os.makedirs(RENDER_OUTPUT_DIR, exist_ok=True)
os.makedirs(TEXTURE_OUTPUT_DIR, exist_ok=True)

# 获取所有皮肤文件
TRAIN_DIR = "train"
skin_files = []
for file in os.listdir(TRAIN_DIR):
    if file.endswith('.png'):
        skin_files.append(os.path.join(TRAIN_DIR, file))

print(f"找到 {len(skin_files)} 个皮肤文件")

async def render_skin(skin_path, output_name):
    """渲染单个皮肤为3D全身视图（正面+背面并排）"""
    try:
        # 生成随机深色背景 (RGB值在0-80之间)
        bg_color = (
            random.randint(0, 80),
            random.randint(0, 80),
            random.randint(0, 80)
        )
        
        # 读取皮肤文件为PIL Image
        skin_img = Image.open(skin_path)
        
        # 创建Skin对象
        skin = Skin(raw_skin=skin_img)
        
        # 为腿脚摆动添加±10度的随机偏移
        leg_offset = random.randint(-10, 10)
        arm_offset = random.randint(-10, 10)
        vr_offset = random.randint(-6, 6)
        hr_offset = random.randint(-20,10)
        ratio_offset = random.randint(-2,0)
        
        # 计算带随机偏移的角度
        #vrll_front = 15 + leg_offset      # 左腿向前
        vrll_front = 15
        #vrrl_front = -15 - leg_offset     # 右腿向后
        vrrl_front = -15
        #vrla_front = -15 - arm_offset     # 左臂向后
        vrla_front = -15
        #vrra_front = 15 + arm_offset      # 右臂向前
        vrra_front = 15
        # vrrl = -15 + vr_offset
        vrrl = -15
        # hrrl = 35 + hr_offset
        hrrl = 35
        # ratiorl = 30 + ratio_offset
        ratiorl = 30

        
        # 渲染正面视图 - 3D行走姿势
        render_front = Render(
            player=skin,
            vr=vrrl,      # 垂直旋转，稍微向下看
            hr=hrrl,       # 水平旋转35度
            hrh=0,       # 头部水平旋转
            vrll=vrll_front,     # 左腿向前（带随机偏移）
            vrrl=vrrl_front,     # 右腿向后（带随机偏移）
            vrla=vrla_front,     # 左臂向后（带随机偏移）
            vrra=vrra_front,     # 右臂向前（带随机偏移）
            ratio=ratiorl,    # 缩放比例
            aa=True,     # 抗锯齿
            display_hair=True,
            display_layers=True
        )
        
        # 渲染背面视图 - 旋转180度（使用相反的角度）
        render_back = Render(
            player=skin,
            vr=vrrl,      # 垂直旋转，稍微向下看
            hr=hrrl + 180,      # 水平旋转215度 (35+180)
            hrh=0,       # 头部水平旋转
            vrll=-vrrl_front,    # 左腿向后（与正面右腿相反）
            vrrl=-vrll_front,    # 右腿向前（与正面左腿相反）
            vrla=-vrra_front,    # 左臂向前（与正面右臂相反）
            vrra=-vrla_front,    # 右臂向后（与正面左臂相反）
            ratio=ratiorl,    # 缩放比例
            aa=True,     # 抗锯齿
            display_hair=True,
            display_layers=True
        )
        
        # 获取渲染结果
        front_img = await render_front.get_render()
        back_img = await render_back.get_render()
        
        # 创建1024x1024的画布，使用随机深色背景
        canvas = Image.new('RGB', (1024, 1024), bg_color)
        
        # 计算每个图像的位置（并排放置）
        front_width, front_height = front_img.size
        back_width, back_height = back_img.size
        
        # 左侧放置正面视图
        front_x = (512 - front_width) // 2
        front_y = (1024 - front_height) // 2
        
        # 右侧放置背面视图
        back_x = 512 + (512 - back_width) // 2
        back_y = (1024 - back_height) // 2
        
        # 粘贴正面视图
        if front_img.mode == 'RGBA':
            canvas.paste(front_img, (front_x, front_y), front_img)
        else:
            canvas.paste(front_img, (front_x, front_y))
        
        # 粘贴背面视图
        if back_img.mode == 'RGBA':
            canvas.paste(back_img, (back_x, back_y), back_img)
        else:
            canvas.paste(back_img, (back_x, back_y))
        
        # 保存渲染图
        render_output_path = os.path.join(RENDER_OUTPUT_DIR, output_name)
        canvas.save(render_output_path)
        
        return True
    except Exception as e:
        print(f"渲染失败 {skin_path}: {str(e)}")
        return False

def scale_texture(skin_path, output_name):
    """缩放原始材质到1024x1024，灰色背景"""
    try:
        # 打开原始皮肤
        img = Image.open(skin_path)
        
        # 处理所有可能有透明度的模式
        # P模式（调色板）、LA模式、RGBA模式都可能有透明度
        if img.mode in ('RGBA', 'LA', 'P'):
            # 先转换为RGBA以保留透明度信息
            if img.mode != 'RGBA':
                img = img.convert('RGBA')
            
            # 创建灰色背景
            background = Image.new('RGB', img.size, (128, 128, 128))
            # 使用alpha通道合成
            background.paste(img, (0, 0), img)
            img = background
        elif img.mode == 'RGB' and 'transparency' in img.info:
            # RGB模式但有透明色信息（通常是某个颜色被指定为透明）
            # 需要转换为RGBA来处理透明度
            # 获取透明色
            trans_color = img.info['transparency']
            
            # 转换为RGBA
            img = img.convert('RGBA')
            
            # 获取图片数据
            datas = img.getdata()
            
            # 将透明色替换为完全透明
            newData = []
            for item in datas:
                # 检查是否是透明色（RGB模式的transparency可能是元组或单个值）
                if isinstance(trans_color, tuple):
                    if item[:3] == trans_color:
                        newData.append((128, 128, 128, 0))  # 完全透明
                    else:
                        newData.append(item)
                else:
                    # 单通道透明度（灰度图）
                    if item[0] == trans_color:
                        newData.append((128, 128, 128, 0))
                    else:
                        newData.append(item)
            
            img.putdata(newData)
            
            # 创建灰色背景并合成
            background = Image.new('RGB', img.size, (128, 128, 128))
            background.paste(img, (0, 0), img)
            img = background
        elif img.mode != 'RGB':
            # 转换其他模式为RGB
            img = img.convert('RGB')
        
        # 使用最近邻插值缩放皮肤材质（保持像素风格）
        # 通常Minecraft皮肤是64x64或64x32
        img_resized = img.resize((1024, 1024), Image.NEAREST)
        
        # 确保保存为RGB模式（没有透明通道）
        if img_resized.mode != 'RGB':
            img_resized = img_resized.convert('RGB')
        
        # 保存缩放后的材质
        texture_output_path = os.path.join(TEXTURE_OUTPUT_DIR, output_name)
        img_resized.save(texture_output_path, 'PNG')
        
        return True
    except Exception as e:
        print(f"缩放失败 {skin_path}: {str(e)}")
        return False

async def process_all_skins():
    """处理所有皮肤"""
    total = len(skin_files)
    success_render = 0
    success_texture = 0
    
    for idx, skin_path in enumerate(skin_files, 1):
        # 获取文件名（不含扩展名）
        filename = os.path.basename(skin_path)
        output_name = filename
        
        print(f"处理 [{idx}/{total}]: {filename}")
        
        # 渲染皮肤
        if await render_skin(skin_path, output_name):
            success_render += 1
        
        # 缩放材质
        if scale_texture(skin_path, output_name):
            success_texture += 1
        
        # 每处理10个文件显示进度
        if idx % 10 == 0:
            print(f"进度: {idx}/{total} - 渲染成功: {success_render}, 材质成功: {success_texture}")
    
    print(f"\n完成!")
    print(f"总计: {total}")
    print(f"渲染成功: {success_render}")
    print(f"材质缩放成功: {success_texture}")
    print(f"渲染图保存在: {RENDER_OUTPUT_DIR}/")
    print(f"缩放材质保存在: {TEXTURE_OUTPUT_DIR}/")

if __name__ == "__main__":
    asyncio.run(process_all_skins())