Convert depth buffer to images

Texture.py

@@ -54,6 +54,30 @@
 X8R8G8B8 = TextureDescription(32, (8, 8, 8), (16, 8, 0))
 
 
+def f16_to_float(f):
+    if f == 0x0:
+        return 0.0
+    f = (f << 11) + 0x3C000000
+    return struct.unpack("f", f.to_bytes(4, "little"))[0]
+
+
+def f24_to_float(f):
+    assert (f >> 24) == 0
+    f &= 0xFFFFFF
+    if f == 0x0:
+        return 0.0
+    f = f << 7
+    return struct.unpack("f", f.to_bytes(4, "little"))[0]
+
+
+ZetaDescription = namedtuple(
+    "ZetaDescription",
+    ["depth_mask", "depth_shift", "stencil_mask", "bpp", "convert_float"],
+)
+Z24S8 = ZetaDescription(0xFFFFFF00, 8, 0xFF, 32, f24_to_float)
+Z16 = ZetaDescription(0xFFFF, 0, 0, 16, f16_to_float)
+
+
 def _decode_texture(
     data, size, pitch, swizzled, bits_per_pixel, channel_sizes, channel_offsets
 ):
@@ -182,6 +206,8 @@ def read_texture_parameters(xbox: Xbox) -> TextureParameters:
 
     swizzle_unk2 = xbox.read_u32(0xFD40086C)
 
+    setup_raster = xbox.read_u32(0xFD401990)
+
     clip_x = (surface_clip_x >> 0) & 0xFFFF
     clip_y = (surface_clip_y >> 0) & 0xFFFF
 
@@ -208,15 +234,18 @@ def read_texture_parameters(xbox: Xbox) -> TextureParameters:
     # FIXME: if surface_type is 0, we probably can't even draw..
 
     format_color = (draw_format >> 12) & 0xF
-    # FIXME: Support 3D surfaces.
-    _format_depth_buffer = (draw_format >> 18) & 0x3
+
+    z_float = (setup_raster >> 29) & 0x1
+    format_depth_buffer = (draw_format >> 18) & 0x3
 
     if not format_color:
         fmt_color = None
     else:
         fmt_color = surface_color_format_to_texture_format(format_color, swizzled)
-    # TODO: Extract swizzle and float state.
-    # fmt_depth = surface_zeta_format_to_texture_format(format_depth_buffer)
+
+    fmt_depth = surface_zeta_format_to_texture_format(
+        format_depth_buffer, swizzled, z_float
+    )
 
     return TextureParameters(
         width=width,
@@ -226,14 +255,101 @@ def read_texture_parameters(xbox: Xbox) -> TextureParameters:
         format_color=fmt_color,
         depth_pitch=depth_pitch,
         depth_offset=depth_offset,
-        format_depth=None,
+        format_depth=fmt_depth,
         surface_type=surface_type,
         swizzle_unk=swizzle_unk,
         swizzle_unk2=swizzle_unk2,
         swizzled=swizzled,
     )
 
 
+def dump_zeta(data, offset, pitch, fmt_depth, width, height):
+    """Convert the zeta buffer at the given offset into a PIL.Image."""
+    depth_img = Image.new("RGB", (width, height))
+    stencil_img = None
+
+    if fmt_depth == 0x2D:
+        info = (True, True, Z16)
+    elif fmt_depth == 0x31:
+        info = (False, True, Z16)
+    elif fmt_depth == 0x2C:
+        info = (True, False, Z16)
+    elif fmt_depth == 0x30:
+        info = (False, False, Z16)
+    elif fmt_depth == 0x2B:
+        info = (True, True, Z24S8)
+        stencil_img = Image.new("RGB", (width, height))
+    elif fmt_depth == 0x2F:
+        info = (False, True, Z24S8)
+        stencil_img = Image.new("RGB", (width, height))
+    elif fmt_depth == 0x2A:
+        info = (True, False, Z24S8)
+        stencil_img = Image.new("RGB", (width, height))
+    elif fmt_depth == 0x2E:
+        info = (False, False, Z24S8)
+        stencil_img = Image.new("RGB", (width, height))
+    else:
+        raise Exception("Unknown depth format :0x%X" % fmt_depth)
+
+    swizzle = info[0]
+    is_float = info[1]
+    fmt = info[2]
+
+    # FIXME: Avoid this nasty ~~convience feature~~ hack
+    if pitch == 0:
+        pitch = width * fmt.bpp // 8
+
+    assert len(data) == pitch * height
+
+    # FIXME: Does this work?
+    if False and swizzle:
+        data = xbox.Unswizzle(data, fmt.bpp, (width, height), pitch)
+
+    depth_values = {}
+    depth_max = 0
+    depth_min = 0xFFFFFF
+    stencil_values = {}
+
+    for y in range(height):
+        for x in range(width):
+            pixel_offset = y * pitch + x * fmt.bpp // 8
+            z_data = data[pixel_offset : pixel_offset + fmt.bpp // 8]
+            zeta = int.from_bytes(z_data, "little")
+            depth = (zeta & fmt.depth_mask) >> fmt.depth_shift
+            stencil = zeta & fmt.stencil_mask
+
+            if is_float:
+                depth = fmt.convert_float(depth)
+
+            depth_values[x, y] = depth
+            stencil_values[x, y] = 255 if stencil != 0 else 0
+
+            if depth > depth_max:
+                depth_max = depth
+
+            if depth < depth_min:
+                depth_min = depth
+
+    # Map depth values to [0, 255]
+    assert depth_max >= depth_min
+    depth_max -= depth_min
+    depth_scaler = 255 / depth_max if depth_max > 0 else 255
+    depth_pixels = depth_img.load()
+    if stencil_img is not None:
+        stencil_pixels = stencil_img.load()
+
+    for y in range(height):
+        for x in range(width):
+            depth = int((depth_values[x, y] - depth_min) * depth_scaler)
+            stencil = stencil_values[x, y]
+
+            depth_pixels[x, y] = (depth, depth, depth)
+            if stencil_img is not None:
+                stencil_pixels[x, y] = (stencil, stencil, stencil)
+
+    return (depth_img, stencil_img)
+
+
 def dump_texture(xbox, offset, pitch, fmt_color, width, height):
     """Convert the texture at the given offset into a PIL.Image."""
     img = None

Trace.py

@@ -420,6 +420,12 @@ def dump_surfaces(self, _data, *_args):
             print("Warning: Invalid color format, skipping surface dump.")
             return []
 
+        if params.depth_offset:
+            depth_buffer = self.xbox.read(
+                Texture.AGP_MEMORY_BASE | params.depth_offset,
+                params.depth_pitch * params.height,
+            )
+
         # Dump stuff we might care about
         self._write("pgraph.bin", _dump_pgraph(self.xbox))
         self._write("pfb.bin", _dump_pfb(self.xbox))
@@ -434,10 +440,7 @@ def dump_surfaces(self, _data, *_args):
         if params.depth_offset and self.enable_raw_pixel_dumping:
             self._write(
                 "mem-3.bin",
-                self.xbox.read(
-                    Texture.AGP_MEMORY_BASE | params.depth_offset,
-                    params.depth_pitch * params.height,
-                ),
+                depth_buffer,
             )
         if self.enable_rdi:
             self._write(
@@ -473,7 +476,6 @@ def dump_surfaces(self, _data, *_args):
         else:
             alpha_path = None
 
-        path = "command%d--color.png" % (self.command_count)
         extra_html = []
 
         extra_html += [img_tags]
@@ -516,6 +518,53 @@ def dump_surfaces(self, _data, *_args):
 
         self._save_image(img, no_alpha_path, alpha_path)
 
+        depth_path = "command%d--depth.png" % (self.command_count)
+        stencil_path = "command%d--stencil.png" % (self.command_count)
+
+        try:
+            if not params.depth_offset:
+                raise Exception("Depth offset is null")
+
+            self._dbg_print("Attempting to dump zeta")
+            depth, stencil = Texture.dump_zeta(
+                depth_buffer,
+                params.depth_offset,
+                params.depth_pitch,
+                params.format_depth,
+                params.width,
+                params.height,
+            )
+        except:
+            depth = None
+            stencil = None
+            print("Failed to dump zeta surface")
+            traceback.print_exc()
+
+        self._save_image(depth, None, depth_path)
+        self._save_image(stencil, None, stencil_path)
+
+        zeta_img_tags = '<img height="128px" src="%s" alt="%s"/>' % (
+            depth_path,
+            depth_path,
+        )
+
+        if stencil is not None:
+            zeta_img_tags += '<img height="128px" src="%s" alt="%s"/>' % (
+                stencil_path,
+                stencil_path,
+            )
+
+        extra_html += [zeta_img_tags]
+        extra_html += [
+            "zeta: [pitch = %d (0x%X)], at 0x%08X, format 0x%X]"
+            % (
+                params.depth_pitch,
+                params.depth_pitch,
+                params.depth_offset,
+                params.format_depth,
+            )
+        ]
+
         return extra_html
 
     def _save_image(self, img, no_alpha_path, alpha_path):