painting_reconstruct.py

import re
from PIL import Image
from pathlib import Path
from typing import Iterable, Optional
from argparse import ArgumentParser

import tkinter as tk
from tkinter import filedialog

import UnityPy
from UnityPy.classes import GameObject
from UnityPy.enums import ClassIDType


VR = re.compile(r'v ')
TR = re.compile(r'vt ')
SR = re.compile(r' ')
def recon(src, mesh):
	sx, sy = src.size
	c = map(SR.split, list(filter(TR.match, mesh))[1::2])
	p = map(SR.split, list(filter(VR.match, mesh))[1::2])
	c = [(round(float(a[1])*sx), round((1-float(a[2]))*sy)) for a in c]
	p = [(-int(float(a[1])), int(float(a[2]))) for a in p]
	my = max(y for x, y in p)
	p = [(x, my-y) for x, y in p[::2]]
	cp = [(l+r, p) for l, r, p in zip(c[::2], c[1::2], p)]
	ox, oy = zip(*[(r-l+p, b-t+q) for (l, t, r, b), (p, q) in cp])
	out = Image.new('RGBA', (max(ox), max(oy)))
	for c, p in cp: out.paste(src.crop(c), p)
	return out

class AzurlaneAsset():
	def __init__(self, offset_dir: Path, asset_sub_path: Path):
		self.directory = offset_dir
		self.bundle = UnityPy.load(str(Path(offset_dir, asset_sub_path)))
		self.container = next(iter(self.bundle.container.values()))

	def loadDependencies(self):
		assetbundle = self.getObjectByPathID(1)
		paths = [str(Path(self.directory, bundlepath)) for bundlepath in assetbundle.m_Dependencies]
		self.bundle.load(paths)

	def getObjectByPathID(self, pathid: int):
		for unity_object in self.bundle.objects:
			if unity_object.path_id == pathid:
				return unity_object.read()

	def getObjectByName(self, name: str, objtype: ClassIDType):
		for unity_object in self.bundle.objects:
			if not unity_object.type == objtype:
				continue
			unity_object = unity_object.read()
			if unity_object.name == name:
				return unity_object

	def getComponentFromObject(self, gameobject: GameObject, types: Iterable[ClassIDType] = None,
		names: set[str] = None, attributes: set[str] = None):

		for component_pptr in gameobject.m_Components:
			if types and not component_pptr.type in types:
				continue

			component = self.getObjectByPathID(component_pptr.path_id)
			if names and not component.name in names:
				continue
			component.read_typetree()
			
			if attributes:
				for attribute in attributes:
					if hasattr(component, attribute):
						return component
				continue
			return component


def parse(x):
	if hasattr(x, "values"):
		return tuple(x.values())
	elif isinstance(x, UnityPy.math.Vector3):
		return (x.X, x.Y, x.Z)
	elif isinstance(x, UnityPy.math.Quaternion):
		return (x.X, x.Y, x.Z, x.W)

class RectTransform:
	def __init__(self, rtf: UnityPy.classes.RectTransform):
		self._obj = rtf
		self.local_rotation = parse(rtf.m_LocalRotation)
		self.local_position = parse(rtf.m_LocalPosition)
		self.local_scale = parse(rtf.m_LocalScale)
		self.anchor_min = parse(rtf.m_AnchorMin)
		self.anchor_max = parse(rtf.m_AnchorMax)
		self.anchored_position = parse(rtf.m_AnchoredPosition)
		self.size_delta = parse(rtf.m_SizeDelta)
		self.pivot = parse(rtf.m_Pivot)


class GameObjectLayer():
	def __init__(self, asset: AzurlaneAsset, gameobject, parent: "GameObjectLayer" = None):
		self.asset = asset
		self.gameobject = gameobject
		self.parent = parent
		self.children = []

		self.transform = asset.getComponentFromObject(gameobject, types=[ClassIDType.RectTransform, ClassIDType.Transform])
		if self.transform.type == ClassIDType.RectTransform:
			self.rect_transform = RectTransform(self.transform)

		self.size = (0, 0)
		self.local_offset = (0, 0)
		self.global_offset = (0, 0)

		self.image = None
		self.meshimage = asset.getComponentFromObject(gameobject, types=[ClassIDType.MonoBehaviour], attributes={"mMesh"})
		if self.meshimage:
			self.loadImage(self.meshimage)

	def __repr__(self):
		return f"<{self.__class__.__name__} {self.gameobject.name}>"

	def loadImage(self, meshimage):
		# reconstruct the image from T2D using Mesh
		self.sprite = self.asset.getObjectByPathID(meshimage.m_Sprite.path_id)
		self.texture2d = self.asset.getObjectByPathID(self.sprite.m_RD.texture.path_id)
		image = self.texture2d.image

		self.mesh = self.asset.getObjectByPathID(meshimage.mMesh.path_id)
		if self.mesh:
			meshexport = self.mesh.export().splitlines()
			image = recon(image, meshexport)

		# recalculate the size of the image (how it is applied inside unity)
		psizex, psizey = image.size
		pdeltax, pdeltay = self.rect_transform.size_delta
		prawx, prawy = meshimage.mRawSpriteSize.values()

		if prawx != psizex or prawy != psizey:
			empty_canvas = Image.new('RGBA', (int(prawx), int(prawy)), (0, 0, 0, 0))
			image = image.transpose(Image.FLIP_TOP_BOTTOM)
			empty_canvas.paste(image)
			image = empty_canvas.transpose(Image.FLIP_TOP_BOTTOM)

		self.image = image.resize((int(pdeltax), int(pdeltay)), Image.LANCZOS)
		self.size = (pdeltax, pdeltay)

	# TODO: CHANGE THIS SHIT
	def loadImageSimple(self, image: Image.Image):
		# recalculate the size of the image (how it is applied inside unity)
		pdeltax, pdeltay = self.rect_transform.size_delta
		#image = image.transpose(Image.FLIP_TOP_BOTTOM)
		self.image = image.resize((int(pdeltax), int(pdeltay)), Image.LANCZOS)
		self.size = (pdeltax, pdeltay)

	def retrieveChildren(self, recursive: bool = True):
		for child in self.transform.m_Children:
			childtf = self.asset.getObjectByPathID(child.path_id)
			#if not isinstance(childtf, OrderedDict): continue # i don't even know what this was for...
			childobject = self.asset.getObjectByPathID(childtf.m_GameObject.path_id)
			objectlayer = GameObjectLayer(self.asset, childobject, self)
			if recursive:
				objectlayer.retrieveChildren(recursive)
				self.children.append(objectlayer)

	def findChildLayer(self, name: str):
		for child in self.children:
			if child.gameobject.name == name:
				return child
			fromchild = child.findChildLayer(name)
			if fromchild:
				return fromchild

	def calculateLocalOffset(self, recursive: bool = True):
		if hasattr(self, 'rect_transform'):
			# calculate relative offset from parent
			anchor_min = self.rect_transform.anchor_min
			anchor_max = self.rect_transform.anchor_max
			anchorpos = self.rect_transform.anchored_position
			size_delta = self.rect_transform.size_delta
			pivot = self.rect_transform.pivot

			if anchor_min == anchor_max:
				# recalculate anchor position for 'special' resized images
				anchorposx = anchorpos[0] + (self.size[0] - size_delta[0]) * pivot[0]
				anchorposy = anchorpos[1] + (self.size[1] - size_delta[1]) * (pivot[1]-1) 
				anchorpos = (anchorposx, anchorposy)

				if self.parent:
					offsetx = self.parent.size[0]*anchor_min[0] + anchorpos[0] - pivot[0]*self.size[0]
					offsety = self.parent.size[1]*anchor_min[1] - anchorpos[1] - (1-pivot[1])*self.size[1]
					self.local_offset = (offsetx, offsety)
			else:
				sizex = (anchor_min[0] - anchor_max[0]) * self.parent.size[0]
				sizey = (anchor_min[1] - anchor_max[1]) * self.parent.size[1]
				self.size = (abs(sizex), abs(sizey))
				self.local_offset = (-anchorpos[0], anchorpos[1])

		if recursive:
			for child in self.children:
				child.calculateLocalOffset(recursive)

	def getSmallestOffset(self):
		min_offx, min_offy = (0, 0)
		if self.image:
			min_offx, min_offy = self.local_offset
		for child in self.children:
			offcx, offcy = child.getSmallestOffset()
			if offcx < min_offx: min_offx = offcx
			if offcy < min_offy: min_offy = offcy
		return min_offx, min_offy

	def calculateGlobalOffset(self, offset = None):
		offset = offset or self.getSmallestOffset()
		self.global_offset = (self.local_offset[0]-offset[0], self.local_offset[1]-offset[1])
		for child in self.children:
			child.calculateGlobalOffset(self.global_offset)

	def getBiggestSize(self):
		sizeoffx, sizeoffy = (0, 0)
		if self.image:
			sizeoffx, sizeoffy = (self.global_offset[0]+self.size[0], self.global_offset[1]+self.size[1])
		for child in self.children:
			csizeoffx, csizeoffy = child.getBiggestSize()
			if csizeoffx > sizeoffx: sizeoffx = csizeoffx
			if csizeoffy > sizeoffy: sizeoffy = csizeoffy
		return sizeoffx, sizeoffy

	def yieldLayers(self):
		if self.image:
			yield self
		for child in self.children:
			for layer in child.yieldLayers():
				yield layer


def get_face_image(facename: str, facetype: str, asset_dir: Path) -> Optional[Image.Image]:
	facebpath = Path("paintingface", facename)
	fasset = AzurlaneAsset(asset_dir, facebpath)
	facet2d = fasset.getObjectByName(facetype, ClassIDType.Texture2D)
	return facet2d.image

def create_image(asset: AzurlaneAsset, append_face: Image.Image = None) -> Image.Image:
	parent_object = asset.getObjectByPathID(asset.container.path_id)
	parent_goi = GameObjectLayer(asset, parent_object)

	parent_goi.retrieveChildren()

	if append_face:
		facelayer = parent_goi.findChildLayer('face')
		facelayer.loadImageSimple(append_face)

	parent_goi.calculateLocalOffset()
	parent_goi.calculateGlobalOffset()

	sizex, sizey = parent_goi.getBiggestSize()
	canvas = Image.new('RGBA', (int(sizex), int(sizey)))
	for layer in parent_goi.yieldLayers():
		#if layer.gameobject.name == "face":
		#	layer.global_offset = (layer.global_offset[0], layer.global_offset[1]+110)
		canvas.alpha_composite(layer.image, (int(layer.global_offset[0]), int(layer.global_offset[1])))
	return canvas

def downscale_image(img: Image.Image) -> Image.Image:
	rx, ry = img.size
	if rx > 2048 or ry > 2048:
		if rx > ry:
			factor = 2048/rx
			return img.resize((2048, round(factor*ry)))
		else:
			factor = 2048/ry
			return img.resize((round(factor*rx), 2048))
	return img


def open_dir_dialog() -> Path:
	root = tk.Tk()
	root.withdraw()
	dir_path = Path(filedialog.askdirectory(initialdir=".", mustexist=True, title="Select asset directory"))
	return dir_path

def open_file_dialog(asset_dir: Path) -> Path:
	root = tk.Tk()
	root.withdraw()
	file_path = Path(filedialog.askopenfilename(initialdir=asset_dir, title="Select painting asset file"))
	return file_path


def main():
	# create argument parser
	parser = ArgumentParser()
	parser.add_argument("-p", "--painting_name", type=str, help="the name of the painting assetbundle file")
	parser.add_argument("-d", "--asset_directory", type=Path, help="directory called 'AssetBundles' containing all client assets")
	parser.add_argument("-f", "--face_name", type=str, default="", help="the name of the face assetbundle file")
	parser.add_argument("-t", "--face_type", type=str, default="0", help="the type/id of the face")
	parser.add_argument("-o", "--out_file", type=Path, help="the output filename")
	args = parser.parse_args()

	# create useful paths
	asset_dir = args.asset_directory
	if asset_dir:
		if asset_dir.name != "AssetBundles":
			raise ValueError("The asset directory needs to be named 'AssetBundles'!")
	else:
		asset_dir = open_dir_dialog()

	painting_name = args.painting_name
	if painting_name:
		painting_subpath = Path("painting", painting_name)
	else:
		painting_fullpath = open_file_dialog(asset_dir)
		painting_subpath = painting_fullpath.relative_to(asset_dir)
	
	targetpath = args.out_file or Path(painting_subpath.name + ".png")

	# load painting asset
	asset = AzurlaneAsset(asset_dir, painting_subpath)
	asset.loadDependencies()

	# get faceimage if required
	face_name = args.face_name
	faceimg = None
	if face_name:
		faceimg = get_face_image(args.face_name, args.face_type, asset_dir)

	# create image, transform it and save
	result = create_image(asset, faceimg)
	result = downscale_image(result)
	result.save(targetpath)


if __name__ == "__main__":
	main()