Hearthstone card art mapping using hearthsim/python-unitypack

cardmap.py

#!/usr/bin/env python
import os
import sys
import glob
import json
import unitypack
from PIL import ImageOps

def get_output_path(filename):
	basedir = "out"
	path = os.path.join(basedir, filename)
	dirs = os.path.dirname(path)
	if not os.path.exists(dirs):
		os.makedirs(dirs)
	return path


def write_to_file(filename, contents, mode="w"):
	path = get_output_path(filename)
	with open(path, mode) as f:
		written = f.write(contents)

	print("Written %i bytes to %r" % (written, path))


def handle_asset(asset, map):
	print(asset)
	count = 0
	for id, obj in asset.objects.items():
		if obj.type == "Texture2D":
			if id in map:
				d = obj.read()
				filename = map[id]["card_id"] + ".png"
				image = d.image
				if image is None:
					print("WARNING: %s is an empty image" % (filename))
					write_to_file(filename, "")
				else:
					img = ImageOps.flip(image)
					path = get_output_path(filename)
					# remove alpha channel and save
					img.convert("RGB").save(path)


def save_as_json(map):
	new_map = {}
	for key, value in map.items():
		new_map[value["card_id"]] = {
			"path_id": key,
			"bundle": value["bundle"],
			"name": value["name"]
		}

	write_to_file("mapping.json", json.dumps(new_map))


def add_mapping(item, map, refs):
	tex_path = item["first"].split("/")
	tex_asset = item["second"]["asset"]
	map[tex_asset.path_id] = {
		"name": tex_path[-1].split(".")[0].lower(),
		"card_id": tex_path[-2].lower(),
		"bundle": refs[tex_asset.file_id].split("/")[-1].lower()
	}


def build_tex_map(dir):
	tex_files = glob.glob(os.path.join(dir, "cardtextures?.unity3d"))
	tex_map = {}

	for file in tex_files:
		with open(file, "rb") as f:
			bundle = unitypack.load(f)

		refs = []
		for asset in bundle.assets:
			# first ref is asset itself
			refs.append(asset.name)
			for r in asset.asset_refs[1:]:
				refs.append(r.file_path)

			for id, obj in asset.objects.items():
				if obj.type == "AssetBundle":
					d = obj.read()
					for item in d["m_Container"]:
						add_mapping(item, tex_map, refs)

	save_as_json(tex_map)

	return tex_map


def extract_card_art(dir, map):
	tex_files = glob.glob(os.path.join(dir, "shared?.unity3d"))  + \
		glob.glob(os.path.join(dir, "cardtextures?.unity3d"))

	for file in tex_files:
		with open(file, "rb") as f:
			bundle = unitypack.load(f)

		for asset in bundle.assets:
			handle_asset(asset, map)


def main():
	hs_dir = sys.argv[1]
	# first pass over files to build mapping
	tex_map = build_tex_map(hs_dir)
	# second pass (includes shared) to extract
	extract_card_art(hs_dir, tex_map)


if __name__ == "__main__":
	main()