| from __future__ import annotations | |
| import argparse,json,time,torch | |
| def chamfer(a,b): | |
| d=torch.cdist(a.float()[None],b.float()[None]).squeeze(0);return float(d.min(1).values.mean()+d.min(0).values.mean()) | |
| def symmetry_x(v): | |
| a=v.float().clone();a[:,0]*=-1;d=torch.cdist(v.float()[None],a[None]).squeeze(0);return float(d.min(1).values.mean()) | |
| def joint_error(a,b):return float(torch.linalg.vector_norm(a.float()-b.float(),dim=-1).mean()) if a.shape==b.shape else None | |
| def main(): | |
| ap=argparse.ArgumentParser();ap.add_argument('--teacher',required=True);ap.add_argument('--student',required=True);ap.add_argument('--out',required=True);a=ap.parse_args();t=torch.load(a.teacher,map_location='cpu',weights_only=False);s=torch.load(a.student,map_location='cpu',weights_only=False) | |
| r={} | |
| if 'vertices' in t and 'vertices' in s:r['chamfer']=chamfer(s['vertices'],t['vertices']);r['symmetry_x']=symmetry_x(s['vertices']) | |
| if 'joints' in t and 'joints' in s:r['joint_error']=joint_error(s['joints'],t['joints']) | |
| if 'runtime_s' in s:r['runtime_s']=s['runtime_s'];r['speedup']=t.get('runtime_s',0)/s['runtime_s'] if t.get('runtime_s') else None | |
| open(a.out,'w').write(json.dumps(r,indent=2));print(json.dumps(r)) | |
| if __name__=='__main__':main() | |