generate_card_textures.py

#!/usr/bin/env python
import json
import os
import sys
import faulthandler; faulthandler.enable()
import UnityPy
from argparse import ArgumentParser
from typing import List, cast, Dict

from PIL import Image, ImageOps
from UnityPy import Environment
from UnityPy.enums import ClassIDType
from UnityPy.helpers import TypeTreeHelper
from UnityPy.classes import PPtr, GameObject, ComponentPair, Tuple, Component, MonoBehaviour, Material, UnityPropertySheet, AssetBundle

class Logger(object):
    def __init__(self, logFile):
        self.terminal = sys.stdout
        self.log = open(logFile, "w")
   
    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)  

    def flush(self):
        # this flush method is needed for python 3 compatibility.
        # this handles the flush command by doing nothing.
        # you might want to specify some extra behavior here.
        pass    

sys.stdout = Logger("generate_card_textures.log")

class CardTextureInfo:
	portrait_path: str
	tile_info: UnityPropertySheet
 
	def __init__(self, portrait_path: str, tile_info: UnityPropertySheet):
		self.portrait_path = portrait_path
		self.tile_info = tile_info

# ./generate_audio.py /e/t > out.txt
# def main():
# 	p = ArgumentParser()
# 	p.add_argument("src")
# 	args = p.parse_args(sys.argv[1:])

# 	sound_effects = extract_info(args.src)
# 	with open('./ref/sound_effects.json', 'w') as resultFile:
# 		resultFile.write(json.dumps(sound_effects))

# ./generate_card_textures.py --outdir out --tiles-dir tiles --skip-existing /e/Games/Hearthstone/Data/Win
def main():
	TypeTreeHelper.read_typetree_c = False

	p = ArgumentParser()
	p.add_argument("src")
	p.add_argument("--outdir", nargs="?", default="out")
	p.add_argument("--skip-existing", action="store_true")
	p.add_argument("--orig-dir", type=str, default="orig", help="Name of output for originals")
	p.add_argument("--tiles-dir", type=str, default="tiles", help="Name of output for tiles")
	p.add_argument(
		"--formats", nargs="*", default=["jpg"],
		help="Which image formats to generate"
	)
	args = p.parse_args(sys.argv[1:])
	generate_card_textures(args.src, args)


def generate_card_textures(src, args): 
	print("Loading environment")
	env: Environment = UnityPy.load(src)
	print("Building cards_map")
	cards_map = build_cards_map(env)
	print("cards_map: %s" % len(cards_map))
	# json.dump(cards_map, sys.stdout, ensure_ascii = False, indent = 4)
	textures_map = build_textures_map(env)
	print("textures_map: %s" % len(textures_map))
	cards_info: Dict[str, CardTextureInfo] = build_cards_info(env, cards_map)
	print("cards_info: %s" % len(cards_info))

	paths = [card.portrait_path for card in cards_info.values()]
	print("Found %i cards, %i textures including %i unique in use." % (
		len(cards_map), len(textures_map), len(set(paths))
	))

	thumb_sizes = (256, 512)
	for card_id, texture_info in cards_info.items():
		try:
			# print("processing %r (%r)" % (texture_info, card_id))
			do_texture(env, card_id, texture_info, textures_map, thumb_sizes, args)
		except Exception as e:
			sys.stderr.write("ERROR on %r (%r): %s\n" % (texture_info, card_id, e))
			raise


def build_cards_map(env: Environment) -> Dict[str, str]:
	for obj in env.objects:
		if obj.type == ClassIDType.MonoBehaviour:
			dataM: MonoBehaviour = cast(MonoBehaviour, obj.read())
			if dataM.m_Name == "cards_map":
				tree = dataM.map
				keys = tree.keys
				values = tree.values
				print("keys: %s" % len(keys))
				print("values: %s" % len(values))
				# Build a dictionary of key => prefabid
				cards_map = {}
				for cardid, value in zip(keys, values):
					# if cardid != "SC_004":
					# 	continue
					# Only keep the id of the prefab, which means what is after prefab:
					asset_id = value.split("prefab:")[1]
					cards_map[cardid] = asset_id
				return cards_map
				
    
def build_textures_map(env: Environment) -> Dict[str, PPtr]:
	textures = {}
	for obj in env.objects:
		if obj.type == ClassIDType.AssetBundle:
			data = cast(AssetBundle, obj.read())
			container = data.m_Container
			for path, asset in container:
				textures[path] = asset.asset
	return textures


def build_cards_info(env: Environment, cards_map: Dict[str, str]):
	cards = {}
	current_card_idx = 0
	# Iterate over the cards map
	for cardid, prefabid in cards_map.items():
		# if cardid != "HERO_02bb":
		# 	continue
		# if current_card_idx < 7760:
		# 	continue
		prefab_pptr = env.container[prefabid]
		print("card %s: %s, prefabid: %s, prefab_info: %s" % (current_card_idx, cardid, prefabid, prefab_pptr))
		current_card_idx += 1
		prefab = cast(GameObject, prefab_pptr.read())
		components = cast(List[ComponentPair], prefab.m_Component)
		# Find the component that is a monobehavior
		for component in components:
			# print("component: %s" % component)
			component_pptr = component.component
			# print("component_pptr: %s" % component_pptr)
			if component_pptr.type.name == "MonoBehaviour":
				card_def = component_pptr.read()
				# print("card_def: %s" % card_def)
				# Sometimes there's multiple per cardid, we remove the ones without art
				if not hasattr(card_def, "m_PortraitTexturePath"):
					print("\tskipping %s, no m_PortraitTexturePath" % cardid)
					continue
				portrait_path = card_def.__getattribute__("m_PortraitTexturePath")
				if ":" in portrait_path:
					portrait_path = portrait_path.split(":")[1]
				if len(portrait_path) == 0:
					print("\tskipping %s, portrait_path is empty" % cardid)
					continue
				# print("portrait_path: %s" % portrait_path)
				tile_ptr: PPtr = card_def.__getattribute__("m_DeckCardBarPortrait")
				# print("tile_ptr: %s" % tile_ptr)
				tile: Material = None if tile_ptr.path_id == 0 else tile_ptr.read()
				# print("tile: %s" % tile)
				tile_info = None if tile == None else tile.m_SavedProperties
				# print("tile_info: %s" % tile_info)
				texture_info: CardTextureInfo = CardTextureInfo(
					portrait_path = portrait_path.lower(),
					tile_info = tile_info,
				)
				cards[cardid] = texture_info

	return cards


def do_texture(env: Environment, card_id: str, texture_info: CardTextureInfo, textures: Dict[str, PPtr], thumb_sizes, args):
	try:
		texture_pptr = textures[texture_info.portrait_path]	
		# print("texture_pptr: %s" % texture_pptr)
		texture = texture_pptr.read()
		# print("texture: %s" % texture)
		flipped = None
		filename, exists = get_filename(args.outdir, args.orig_dir, card_id, ext=".png")

		if not (args.skip_existing and exists):
			print("-> %r" % (filename))
			flipped = ImageOps.scale(texture.image, 1).convert("RGB")
			flipped.save(filename)

		for format in args.formats:
			ext = "." + format
			filename, exists = get_filename(args.outdir, args.tiles_dir, card_id, ext=ext)
			if not (args.skip_existing and exists):
				print("will build texture for %r" % (filename))
				if not texture.image:
					print("texture has no image")
				tile_texture = generate_tile_image(env, texture.image, texture_info.tile_info)
				if not tile_texture:
					print("could not generate tile texture %s" % card_id)
					# Some hero skins have no tiles, but still have the thumb
				else:
					print("-> %r" % (filename))
					tile_texture.save(filename)
					

			if ext == ".png":
				# skip png generation for thumbnails
				continue

			for sz in thumb_sizes:
				thumb_dir = "%ix" % (sz)
				filename, exists = get_filename(args.outdir, thumb_dir, card_id, ext=ext)
				if not (args.skip_existing and exists):
					if not flipped:
						flipped = ImageOps.scale(texture.image, 1).convert("RGB")
					thumb_texture = flipped.resize((sz, sz))
					print("-> %r" % (filename))
					thumb_texture.save(filename)
	except Exception as e:
		print("ERROR on %r (%r): %s" % (texture_info, card_id, e))
        


def generate_tile_image(env: Environment, img, tile_info):
	if tile_info is None:
		return None
    
	if (img.width, img.height) != (512, 512):
		img = img.resize((512, 512), Image.ANTIALIAS)

	# tile the image horizontally (x2 is enough),
	# some cards need to wrap around to create a bar (e.g. Muster for Battle),
	# also discard alpha channel (e.g. Soulfire, Mortal Coil)
	tiled = Image.new("RGB", (img.width * 2, img.height))
	tiled.paste(img, (0, 0))
	tiled.paste(img, (img.width, 0))

	# print("tile: %s" % tile_info)
	# props = (-0.13, 0.13, 1, 1, 0, 0, 1.1, img.width)
	props = (-0.2, 0.25, 1, 1, 0, 0, 1, img.width)
	# TODO: fix this
	# if tile:
	# 	print("texEnvs: %s" % tile.m_TexEnvs)
	# 	print("texEnvs 2: %s" % tile.m_TexEnvs[''])
	# 	texEnvs = tile.m_TexEnvs['']
	# 	m_Texture: PPtr = texEnvs.m_Texture
	# 	tex = m_Texture.read()
	# 	print("m_Texture: %s" % m_Texture)
	# 	m_Offset = texEnvs.m_Offset
	# 	print("m_Offset: %s" % m_Offset)
	# 	m_Scale = texEnvs.m_Scale
	# 	print("m_Scale: %s" % m_Scale)
	# 	m_Floats = tile.m_Floats
	# 	print("m_Floats: %s" % m_Floats)
	# 	print("tex: %s" % tex)
	# 	props = (
	# 		tile.m_TexEnvs["_MainTex"].m_Offset.X,
	# 		tile.m_TexEnvs["_MainTex"].m_Offset.Y,
	# 		tile.m_TexEnvs["_MainTex"].m_Scale.X,
	# 		tile.m_TexEnvs["_MainTex"].m_Scale.Y,
	# 		tile.m_Floats.get("_OffsetX", 0.0),
	# 		tile.m_Floats.get("_OffsetY", 0.0),
	# 		tile.m_Floats.get("_Scale", 1.0),
	# 		img.width
	# 	)
		# print("tile props: %s" % (props,))

	x, y, width, height = get_rect(*props)
	# print("x: %s, y: %s, width: %s, height: %s" % (x, y, width, height))
	bar = ImageOps.flip(tiled).crop((x, y, x + width, y + height))
	bar = ImageOps.flip(bar)
	# negative x scale means horizontal flip
	if props[2] < 0:
		bar = ImageOps.mirror(bar)

	return bar.resize((OUT_WIDTH, OUT_HEIGHT), Image.LANCZOS)


# Deck tile generation
TEX_COORDS = [(0.0, 0.3856), (1.0, 0.6144)]
OUT_DIM = 256
OUT_WIDTH = round(TEX_COORDS[1][0] * OUT_DIM - TEX_COORDS[0][0] * OUT_DIM)
OUT_HEIGHT = round(TEX_COORDS[1][1] * OUT_DIM - TEX_COORDS[0][1] * OUT_DIM)


def get_rect(ux, uy, usx, usy, sx, sy, ss, tex_dim=512):
	# calc the coords
	tl_x = ((TEX_COORDS[0][0] + sx) * ss) * usx + ux
	tl_y = ((TEX_COORDS[0][1] + sy) * ss) * usy + uy
	br_x = ((TEX_COORDS[1][0] + sx) * ss) * usx + ux
	br_y = ((TEX_COORDS[1][1] + sy) * ss) * usy + uy

	# adjust if x coords cross-over
	horiz_delta = tl_x - br_x
	if horiz_delta > 0:
		tl_x -= horiz_delta
		br_x += horiz_delta

	# get the bar rectangle at tex_dim size
	x = round(tl_x * tex_dim)
	y = round(tl_y * tex_dim)
	width = round(abs((br_x - tl_x) * tex_dim))
	height = round(abs((br_y - tl_y) * tex_dim))

	# adjust x and y, so that texture is "visible"
	x = (x + width) % tex_dim - width
	y = (y + height) % tex_dim - height

	# ??? to cater for some special cases
	min_visible = tex_dim / 4
	while x + width < min_visible:
		x += tex_dim
	while y + height < 0:
		y += tex_dim

	# ensure wrap around is used
	if x < 0:
		x += tex_dim

	return (x, y, width, height)


def get_dir(basedir, dirname):
	ret = os.path.join(basedir, dirname)
	if not os.path.exists(ret):
		os.makedirs(ret)
	return ret


def get_filename(basedir, dirname, name, ext=".png"):
	dirname = get_dir(basedir, dirname)
	filename = name + ext
	path = os.path.join(dirname, filename)
	return path, os.path.exists(path)



if __name__ == "__main__":
	main()