retroDE_ps2/tools/gs_sh3_frame_ref.py

#!/usr/bin/env python3
"""retroDE_ps2 — Ch349 step 4 (definitive): composite a WHOLE SH3 frame from GS local memory.

Per-draw reconstruction (gs_sh3_draw_ref) proves one surface; this composites EVERY textured draw of a frame
to pixel-check the reconstruction against the real PCSX2 screenshot — the strongest faithfulness test. It walks
draws in capture order, keeping a GS local-memory model live (replaying host->local uploads as their idx is
passed, so each draw samples the texture state IT saw, not a stale end-of-frame one), reconstructs+caches each
bound texture (PSMT8 via grid-CSM1 CLUT, or PSMCT32 direct), and rasterizes perspective-correct with a z-buffer.

Output: composited frame PNG + a side-by-side vs the screenshot + a coverage/color summary. SH3-derived ->
LOCAL/gitignored.
Usage:  gs_sh3_frame_ref.py <dump.gs.zst> [--frame F] [--shot path.png] [--max-draws N] [--out DIR]
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gs_sh3_draw_census as C
import gs_sh3_recon as RC
import gs_localmem as LM
import gs_texture_residency as R

SW, SH = 640, 480

def edge(ax, ay, bx, by, px, py):
    return (px-ax)*(by-ay) - (py-ay)*(bx-ax)

def reconstruct_texture(mem, t0):
    """Return (idx_or_words, pal_or_None, kind) for the bound texture. PSMT8 -> (idx, pal,'i8'); PSMCT32 ->
    (words, None,'ct32')."""
    tbp, psm, tbw, tw, th = t0["tbp"], t0["psm"], t0["tbw"], t0["tw"], t0["th"]
    if psm == 0x13:   # PSMT8 indexed
        idx = mem.read_psmt8(tbp, tbw, tw, th)
        pal = RC.read_clut32(mem, t0["cbp"], order="grid")
        return idx, pal, "i8"
    if psm == 0x00:   # PSMCT32 direct
        words = [mem.read_ct32_word(tbp, tbw, x, y) for y in range(th) for x in range(tw)]
        return words, None, "ct32"
    return None, None, None

def sample(tex, pal, kind, tw, th, u, v):
    tx = int(u) % tw; ty = int(v) % th
    if tx < 0: tx += tw
    if ty < 0: ty += th
    if kind == "i8":
        p = pal[tex[ty*tw+tx] & 0xFF]
    else:
        p = tex[ty*tw+tx]
    return (p&0xFF, (p>>8)&0xFF, (p>>16)&0xFF)

def main(argv):
    if len(argv) < 2:
        print(__doc__); return 2
    dump = argv[1]
    def opt(n, dv=None): return argv[argv.index(n)+1] if n in argv else dv
    frame = int(opt("--frame","1")); maxd = int(opt("--max-draws","100000"))
    min_area = float(opt("--min-area","0"))
    outdir = opt("--out", os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                                       "captures","gs","silenthill3","extracted","recon"))
    os.makedirs(outdir, exist_ok=True)

    # all draws of the frame, in capture (idx) order, with full geometry (parse ONCE, reuse for census)
    collected = R.collect(dump, 0)
    d, h, events, uploads, runs, vram = collected
    draws, _, _ = C.census(dump, frame_filter=frame, min_prims=1, collected=collected)
    draws = [dr for dr in draws if dr["prim"]["tme"] and dr["tex0"]["psm"] in (0x13,0x00)
             and dr["onscreen_area"] >= min_area]
    draws.sort(key=lambda x: x["first_idx"])
    print(f"[frameref] frame {frame}: {len(draws)} textured (PSMT8/CT32) draws, min_area={min_area}; "
          f"replaying uploads incrementally")

    mem = LM.LocalMem(vram)
    up_sorted = sorted([u for u in uploads if u["dpsm"]==0x00], key=lambda u: u["idx"])
    up_i = 0
    # epoch counter per tbp (bumped whenever we apply an upload to that dbp) -> texture cache key
    epoch = {}
    texcache = {}
    fb = [(0,0,0,0)]*(SW*SH); zb = [-1]*(SW*SH)

    def apply_uploads_before(idx):
        nonlocal up_i
        while up_i < len(up_sorted) and up_sorted[up_i]["idx"] < idx:
            u = up_sorted[up_i]
            off, end = u["blob_range"]; blob = d[off:end]
            words = [int.from_bytes(blob[i*4:i*4+4],"little") for i in range(len(blob)//4)]
            mem.write_image_ct32(u["dbp"], u["dbw"], u["dx"], u["dy"], u["w"], u["h"], words)
            epoch[u["dbp"]] = epoch.get(u["dbp"],0)+1
            up_i += 1

    painted_total = 0
    for n, dr in enumerate(draws[:maxd]):
        apply_uploads_before(dr["first_idx"])
        t0 = dr["tex0"]
        key = (t0["tbp"], t0["psm"], t0["tbw"], t0["tw"], t0["th"], t0.get("cbp"),
               epoch.get(t0["tbp"],0), epoch.get(t0.get("cbp"),0))
        if key not in texcache:
            texcache[key] = reconstruct_texture(mem, t0)
        tex, pal, kind = texcache[key]
        if kind is None: continue
        tw, th = t0["tw"], t0["th"]
        verts = dr["verts"]; pt = dr["prim"]["type"]
        def tri(i0,i1,i2):
            nonlocal painted_total
            v0,v1,v2 = verts[i0],verts[i1],verts[i2]
            x0,y0=v0["x"],v0["y"]; x1,y1=v1["x"],v1["y"]; x2,y2=v2["x"],v2["y"]
            minx=max(0,int(min(x0,x1,x2))); maxx=min(SW-1,int(max(x0,x1,x2))+1)
            miny=max(0,int(min(y0,y1,y2))); maxy=min(SH-1,int(max(y0,y1,y2))+1)
            if maxx<minx or maxy<miny: return
            area=edge(x0,y0,x1,y1,x2,y2)
            if abs(area)<1e-6: return
            inv=1.0/area
            for py in range(miny,maxy+1):
                base=py*SW
                for px in range(minx,maxx+1):
                    cx,cy=px+0.5,py+0.5
                    w0=edge(x1,y1,x2,y2,cx,cy); w1=edge(x2,y2,x0,y0,cx,cy); w2=edge(x0,y0,x1,y1,cx,cy)
                    if not ((w0>=0 and w1>=0 and w2>=0) or (w0<=0 and w1<=0 and w2<=0)): continue
                    b0,b1,b2=w0*inv,w1*inv,w2*inv
                    Q=b0*v0["q"]+b1*v1["q"]+b2*v2["q"]
                    if abs(Q)<1e-12: continue
                    S=b0*v0["s"]+b1*v1["s"]+b2*v2["s"]; T=b0*v0["t"]+b1*v1["t"]+b2*v2["t"]
                    u=(S/Q)*tw; vv=(T/Q)*th
                    z=int(b0*v0["z"]+b1*v1["z"]+b2*v2["z"])
                    o=base+px
                    if z>=zb[o]:
                        zb[o]=z; fb[o]=sample(tex,pal,kind,tw,th,u,vv)+(255,); painted_total+=1
        if pt==4:
            for i in range(2,len(verts)): tri(i-2,i-1,i)
        elif pt==5:
            for i in range(2,len(verts)): tri(0,i-1,i)
        elif pt==3:
            for i in range(0,len(verts)-2,3): tri(i,i+1,i+2)
        if (n+1)%200==0: print(f"   ...{n+1}/{len(draws)} draws, {painted_total} px, {len(texcache)} tex cached")

    cov = sum(1 for p in fb if p[3])
    RC.save_png(os.path.join(outdir, f"frame{frame}_composite.png"), SW, SH, fb)
    print(f"[frameref] composite: {cov}/{SW*SH} px painted ({100*cov/(SW*SH):.1f}%), {len(texcache)} textures")
    # side-by-side vs screenshot
    shot = opt("--shot")
    if shot and os.path.exists(shot):
        from PIL import Image
        comp = Image.open(os.path.join(outdir, f"frame{frame}_composite.png")).convert("RGB")
        gt = Image.open(shot).convert("RGB").resize((SW,SH))
        sbs = Image.new("RGB",(SW*2+8,SH),(40,40,40))
        sbs.paste(comp,(0,0)); sbs.paste(gt,(SW+8,0))
        sbs.save(os.path.join(outdir, f"frame{frame}_vs_screenshot.png"))
        print(f"[frameref] wrote frame{frame}_vs_screenshot.png (left=recon, right=PCSX2)")
    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv))