import cv2,numpy as np,sys,json
# track a similarity transform of the upper body from a seed frame forwards and backwards
d=sys.argv[1]; seed=int(sys.argv[2]); a0,a1=int(sys.argv[3]),int(sys.argv[4])
cx,cy,R=map(float,sys.argv[5:8])   # seed: region centre + radius (upper body)
out=sys.argv[8]
G=lambda i: cv2.cvtColor(cv2.imread(f'{d}/f{i:03d}.png'),cv2.COLOR_BGR2GRAY)
T={seed:np.eye(3)}   # maps seed coords -> frame coords
def run(rng):
    prev=seed
    for i in rng:
        g0,g1=G(prev),G(i); M=T[prev]
        c=M@[cx,cy,1]; s=np.sqrt(abs(np.linalg.det(M[:2,:2]))); r=R*s
        mk=np.zeros_like(g0); cv2.circle(mk,(int(c[0]),int(c[1])),int(r),255,-1)
        p0=cv2.goodFeaturesToTrack(g0,400,0.01,5,mask=mk)
        p1,st,_=cv2.calcOpticalFlowPyrLK(g0,g1,p0,None,winSize=(31,31),maxLevel=4)
        pb,st2,_=cv2.calcOpticalFlowPyrLK(g1,g0,p1,None,winSize=(31,31),maxLevel=4)
        ok=(st[:,0]==1)&(st2[:,0]==1)&(np.linalg.norm(pb-p0,axis=2)[:,0]<1.0)
        if ok.sum()<8: print("lost",i,file=sys.stderr); break
        A,inl=cv2.estimateAffinePartial2D(p0[ok],p1[ok],method=cv2.RANSAC,ransacReprojThreshold=2.0)
        A3=np.vstack([A,[0,0,1]]); T[i]=A3@M
        print(i,ok.sum(),int(inl.sum()),file=sys.stderr); prev=i
run(range(seed+1,a1+1)); run(range(seed-1,a0-1,-1))
json.dump({str(k):v.tolist() for k,v in T.items()},open(out,'w'))