Upload generate_dataset.py

generate_dataset.py

@@ -0,0 +1,242 @@
+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())