139 lines
5.4 KiB
Python
139 lines
5.4 KiB
Python
import os
|
|
import requests
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont, ImageEnhance
|
|
from pathlib import Path
|
|
from io import BytesIO
|
|
import itertools
|
|
|
|
class ChessAssetProcessor:
|
|
def __init__(self):
|
|
self.base_dir = Path.cwd()
|
|
self.raw_dir = self.base_dir / 'raw'
|
|
self.training_dir = self.base_dir / 'training'
|
|
|
|
# URLs for pieces
|
|
self.piece_base_url = 'https://www.chess.com/chess-themes/pieces/neo/300'
|
|
|
|
# Piece configurations - matches the Swift enum exactly
|
|
self.pieces = {
|
|
'white': ['pawn', 'knight', 'bishop', 'rook', 'queen', 'king'],
|
|
'black': ['pawn', 'knight', 'bishop', 'rook', 'queen', 'king']
|
|
}
|
|
|
|
# Square colors - used for background variations only
|
|
self.square_colors = {
|
|
'light': '#eeeed2',
|
|
'dark': '#759656',
|
|
'light_highlighted': '#f6f68d',
|
|
'dark_highlighted': '#bdcc49'
|
|
}
|
|
|
|
# Border configurations
|
|
self.border_configs = [
|
|
{}, # No borders
|
|
{'top': True},
|
|
{'bottom': True},
|
|
{'left': True},
|
|
{'top': True, 'left': True},
|
|
{'bottom': True, 'left': True},
|
|
]
|
|
|
|
self._setup_directories()
|
|
|
|
def _setup_directories(self):
|
|
"""Create directory structure for training data"""
|
|
(self.raw_dir / 'pieces').mkdir(parents=True, exist_ok=True)
|
|
self.training_dir.mkdir(exist_ok=True)
|
|
|
|
# Create directories for each piece type
|
|
for color in ['white', 'black']:
|
|
for piece in self.pieces[color]:
|
|
(self.training_dir / f"{color}_{piece}").mkdir(parents=True, exist_ok=True)
|
|
|
|
def hex_to_rgb(self, hex_color):
|
|
"""Convert hex color to RGB tuple"""
|
|
hex_color = hex_color.lstrip('#')
|
|
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
|
|
|
def create_base_square(self, color, size=100):
|
|
"""Create a square with specified color"""
|
|
rgb_color = self.hex_to_rgb(color)
|
|
return Image.new('RGB', (size, size), rgb_color)
|
|
|
|
def add_borders(self, image, borders, border_color=(0, 0, 0)):
|
|
"""Add borders according to configuration"""
|
|
w, h = image.size
|
|
result = image.copy()
|
|
draw = ImageDraw.Draw(result)
|
|
border_size = 5
|
|
|
|
if borders.get('top'):
|
|
draw.line([(0, 0), (w-1, 0)], fill=border_color, width=border_size)
|
|
if borders.get('bottom'):
|
|
draw.line([(0, h-1), (w-1, h-1)], fill=border_color, width=border_size)
|
|
if borders.get('left'):
|
|
draw.line([(0, 0), (0, h-1)], fill=border_color, width=border_size)
|
|
|
|
return result
|
|
|
|
def download_piece(self, color, piece, target_size):
|
|
"""Download a specific chess piece"""
|
|
piece_letter = piece[0] if piece != 'knight' else 'n'
|
|
piece_url = f"{self.piece_base_url}/{color[0]}{piece_letter}.png"
|
|
|
|
response = requests.get(piece_url)
|
|
if response.status_code == 200:
|
|
piece_img = Image.open(BytesIO(response.content)).convert('RGBA')
|
|
return piece_img.resize((target_size, target_size), Image.Resampling.LANCZOS)
|
|
return None
|
|
|
|
def create_training_data(self):
|
|
"""Create comprehensive training dataset"""
|
|
square_size = 100
|
|
|
|
print("Downloading pieces and creating variations...")
|
|
for color in ['white', 'black']:
|
|
for piece in self.pieces[color]:
|
|
piece_dir = self.training_dir / f"{color}_{piece}"
|
|
print(f"\nProcessing {color} {piece}...")
|
|
|
|
# Download piece
|
|
piece_img = self.download_piece(color, piece, square_size)
|
|
if piece_img is None:
|
|
print(f"Failed to download {color} {piece}")
|
|
continue
|
|
|
|
# Create variations with different backgrounds
|
|
variation_count = 0
|
|
for bg_name, bg_color in self.square_colors.items():
|
|
# Create base square
|
|
base_square = self.create_base_square(bg_color, square_size)
|
|
|
|
# Add border variations
|
|
for border_config in self.border_configs:
|
|
# Add borders
|
|
bordered = self.add_borders(base_square, border_config)
|
|
|
|
# Convert to RGBA for composition
|
|
bordered_rgba = bordered.convert('RGBA')
|
|
|
|
# Combine with piece
|
|
combined = Image.alpha_composite(bordered_rgba, piece_img)
|
|
|
|
# Generate filename
|
|
border_desc = '_'.join(k for k,v in border_config.items() if v)
|
|
filename = f"{bg_name}_{border_desc}_{variation_count}.png" if border_desc else f"{bg_name}_{variation_count}.png"
|
|
|
|
# Save image
|
|
combined.save(piece_dir / filename)
|
|
variation_count += 1
|
|
|
|
print(f"Created {variation_count} variations for {color} {piece}")
|
|
|
|
def main():
|
|
processor = ChessAssetProcessor()
|
|
processor.create_training_data()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|