Add support for rendering only certain depth planes for debugging.

bemani/format/afp/render.py

@@ -139,7 +139,7 @@ class AFPRenderer(VerboseOutput):
             data.parse()
         self.swfs[name] = data
 
-    def render_path(self, path: str, background_color: Optional[Color] = None, verbose: bool = False) -> Tuple[int, List[Image.Image]]:
+    def render_path(self, path: str, background_color: Optional[Color] = None, verbose: bool = False, only_depths: Optional[List[int]] = None) -> Tuple[int, List[Image.Image]]:
         # Given a path to a SWF root animation or an exported animation inside a SWF,
         # attempt to render it to a list of frames, one per image.
         components = path.split(".")
@@ -152,7 +152,7 @@ class AFPRenderer(VerboseOutput):
                 # This is the SWF we care about.
                 with self.debugging(verbose):
                     swf.color = background_color or swf.color
-                    return self.__render(swf, components[1] if len(components) > 1 else None)
+                    return self.__render(swf, components[1] if len(components) > 1 else None, only_depths=only_depths)
 
         raise Exception(f'{path} not found in registered SWFs!')
 
@@ -366,7 +366,15 @@ class AFPRenderer(VerboseOutput):
         else:
             raise Exception(f"Failed to process tag: {tag}")
 
-    def __render_object(self, img: Image.Image, renderable: PlacedObject, parent_transform: Matrix, parent_origin: Point, prefix: str="") -> Image.Image:
+    def __render_object(
+        self,
+        img: Image.Image,
+        renderable: PlacedObject,
+        parent_transform: Matrix,
+        parent_origin: Point,
+        only_depths: Optional[List[int]] = None,
+        prefix: str="",
+    ) -> Image.Image:
         self.vprint(f"{prefix}  Rendering placed object ID {renderable.object_id} from sprite {renderable.source.tag_id} onto Depth {renderable.depth}")
 
         # Compute the affine transformation matrix for this object.
@@ -390,7 +398,7 @@ class AFPRenderer(VerboseOutput):
                 key=lambda obj: obj.depth,
             )
             for obj in objs:
-                img = self.__render_object(img, obj, transform, parent_origin.add(renderable.rotation_offset), prefix=prefix + " ")
+                img = self.__render_object(img, obj, transform, parent_origin.add(renderable.rotation_offset), only_depths=only_depths, prefix=prefix + " ")
         elif isinstance(renderable, PlacedShape):
             # This is a shape draw reference.
             shape = renderable.source
@@ -405,6 +413,9 @@ class AFPRenderer(VerboseOutput):
                 if not (params.flags & 0x1):
                     # Not instantiable, don't render.
                     return img
+                if only_depths is not None and renderable.depth not in only_depths:
+                    # Not on the correct depth plane.
+                    return img
 
                 if params.flags & 0x8:
                     # TODO: Need to support blending and UV coordinate colors here.
@@ -507,7 +518,7 @@ class AFPRenderer(VerboseOutput):
 
         return None
 
-    def __render(self, swf: SWF, export_tag: Optional[str]) -> Tuple[int, List[Image.Image]]:
+    def __render(self, swf: SWF, export_tag: Optional[str], only_depths: Optional[List[int]] = None) -> Tuple[int, List[Image.Image]]:
         # If we are rendering an exported tag, we want to perform the actions of the
         # rest of the SWF but not update any layers as a result.
         visible_tag = None
@@ -562,7 +573,7 @@ class AFPRenderer(VerboseOutput):
                     clip = self.__find_renderable(root_clip, visible_tag)
                     if clip:
                         for obj in sorted(clip.placed_objects, key=lambda obj: obj.depth):
-                            curimage = self.__render_object(curimage, obj, root_clip.transform, root_clip.rotation_offset)
+                            curimage = self.__render_object(curimage, obj, root_clip.transform, root_clip.rotation_offset, only_depths=only_depths)
                 else:
                     # Nothing changed, make a copy of the previous render.
                     self.vprint("  Using previous frame render")

bemani/utils/afputils.py

@@ -267,6 +267,12 @@ def main() -> int:
         default=None,
         help="Set the background color of the animation as a comma-separated RGB or RGBA color, overriding a default if present in the SWF.",
     )
+    render_parser.add_argument(
+        "--only-depths",
+        type=str,
+        default=None,
+        help="Only render objects on these depth planes. Can provide either a number or a comma-separated list of numbers.",
+    )
     render_parser.add_argument(
         "--disable-threads",
         action="store_true",
@@ -293,7 +299,6 @@ def main() -> int:
         help="Disable multi-threaded rendering.",
     )
 
-
     args = parser.parse_args()
 
     if args.action == "extract":
@@ -710,7 +715,11 @@ def main() -> int:
                 color = None
 
             # Render the gif/webp frames.
-            duration, images = renderer.render_path(args.path, verbose=args.verbose, background_color=color)
+            if args.only_depths is not None:
+                depths = [int(d.strip()) for d in args.only_depths.split(",")]
+            else:
+                depths = None
+            duration, images = renderer.render_path(args.path, verbose=args.verbose, background_color=color, only_depths=depths)
 
             if len(images) == 0:
                 raise Exception("Did not render any frames!")