Get animations script to work with normal bootlegs

cbpickaxe/__init__.py

@@ -2,6 +2,7 @@
 A library for data mining the game Cassette Beasts.
 """
 from .animation import Animation, Frame, FrameTag, Box
+from .elemental_type import ElementalType
 from .hoylake import Hoylake
 from .item import Item
 from .misc_types import Color
@@ -12,6 +13,7 @@
 __all__ = [
     "Animation",
     "Box",
+    "ElementalType",
     "Frame",
     "FrameTag",
     "Hoylake",

cbpickaxe/elemental_type.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass
+from typing import cast, IO, List
+
+import godot_parser as gp
+
+from .misc_types import Color
+
+
+@dataclass
+class ElementalType:
+    palette: List[Color]
+
+    @staticmethod
+    def from_tres(input_stream: IO[str]) -> "ElementalType":
+        scene = gp.parse(input_stream.read())
+
+        palette = None
+
+        for section in scene.get_sections():
+            # pylint: disable-next=unidiomatic-typecheck
+            if type(section) == gp.GDResourceSection:
+                palette = section["palette"]
+
+        assert isinstance(palette, list)
+
+        for color in palette:
+            assert isinstance(color, gp.Color)
+        palette = cast(List[gp.Color], palette)
+
+        return ElementalType(palette=[Color.from_gp(color) for color in palette])

cbpickaxe/hoylake.py

@@ -11,6 +11,7 @@
 import re
 
 from .animation import Animation
+from .elemental_type import ElementalType
 from .item import Item
 from .monster_form import MonsterForm
 from .move import Move
@@ -60,6 +61,21 @@ def load_root(self, name: str, new_root: str | os.PathLike) -> None:
         self.__roots[name] = new_root
         self.__load_translation_tables(new_root)
 
+    def load_elemental_type(self, path: str) -> Tuple[RootName, ElementalType]:
+        self.__check_if_root_loaded()
+
+        relative_path = Hoylake.__parse_res_path(path)
+
+        for root_name, root in self.__roots.items():
+            type_path = root / relative_path
+            if type_path.exists():
+                with open(type_path, "r", encoding="utf-8") as input_stream:
+                    elemental_type = ElementalType.from_tres(input_stream)
+
+                    return root_name, elemental_type
+
+        raise ValueError(f"Could not find elemental type file at path: {path}")
+
     def load_animation(self, path: str) -> Animation:
         """
         Loads in the animation at the given res:// filepath.

cbpickaxe/misc_types.py

@@ -5,6 +5,7 @@
 write out to JSON.
 """
 from dataclasses import dataclass
+from typing import Tuple
 
 import godot_parser as gp
 
@@ -22,6 +23,17 @@ class Color:
     blue: float  #: Blue component in the range of [0.0, 1.0].
     alpha: float  #: Alpha/opacity of the color in the range of [0.0, 1.0], where 1.0 indicates fully opaque and 0.0 indicates fully transparent.
 
+    def to_8bit_rgba(self) -> Tuple[int, int, int, int]:
+        """
+        Converts the color to an RGBA tuple where each color value is in the range [0, 255].
+        """
+        return (
+            int(round(self.red * 255)),
+            int(round(self.green * 255)),
+            int(round(self.blue * 255)),
+            int(round(self.alpha * 255)),
+        )
+
     @staticmethod
     def from_gp(original: gp.Color) -> "Color":
         """

cbpickaxe_scripts/generate_monster_animations.py

@@ -29,13 +29,18 @@ def main(argv: List[str]) -> int:
     )
     parser.add_argument("--output_directory", required=True)
     parser.add_argument("--crop", default=False, action="store_true")
+    parser.add_argument("--bootleg_type", default=None)
 
     args = parser.parse_args(argv)
 
     hoylake = cbp.Hoylake()
     for i, root in enumerate(args.roots):
         hoylake.load_root(str(i), pathlib.Path(root))
 
+    bootleg_type = None
+    if args.bootleg_type is not None:
+        _, bootleg_type = hoylake.load_elemental_type(args.bootleg_type)
+
     monsters = {}
     for monsters_path in args.monster_form_paths:
         if monsters_path.endswith(".tres"):
@@ -64,6 +69,8 @@ def main(argv: List[str]) -> int:
 
         image_filepath = hoylake.lookup_filepath(image_filepath_relative)
         source_image = PIL.Image.open(image_filepath)
+        if bootleg_type is not None:
+            source_image = recolor_to_bootleg(source_image, monster_form, bootleg_type)
 
         monster_name = hoylake.translate(monster_form.name)
         seen_monster_names[monster_name] += 1
@@ -117,5 +124,42 @@ def main(argv: List[str]) -> int:
     return SUCCESS
 
 
+def recolor_to_bootleg(
+    image: PIL.Image.Image,
+    monster_form: cbp.MonsterForm,
+    elemental_type: cbp.ElementalType,
+) -> PIL.Image.Image:
+    if len(monster_form.swap_colors) < 5:
+        print(
+            f"Warning: Insufficient swap colors for monster_form: {monster_form.name}"
+        )
+        return image
+
+    assert (
+        len(elemental_type.palette) >= 5
+    ), f"Elemental type's palette only has {len(elemental_type.palette)} colors. Must be at least 5."
+
+    color_mapping = {
+        monster_form.swap_colors[i]
+        .to_8bit_rgba(): elemental_type.palette[i]
+        .to_8bit_rgba()
+        for i in range(0, 5)
+    }
+
+    # This appears to be the correct way to do it in Pillow. The point method appears not to
+    # support non-greyscale images.
+    new_image = image.copy()
+    pixels = new_image.load()
+    colors = set()
+    for i in range(new_image.size[0]):
+        for j in range(new_image.size[1]):
+            new_color = color_mapping.get(pixels[i, j], None)
+            colors.add(pixels[i, j])
+            if new_color is not None:
+                pixels[i, j] = new_color
+
+    return new_image
+
+
 def main_without_args() -> int:
     return main(sys.argv[1:])