hUddlZdZdZdZy)Nc |jdddd}tt|ddddddddf|}|t|dddddddf||d<|S) zB extract 3D points/confidence from prediction head output rN)mode)pts3dconf)permutedictreg_dense_depthreg_dense_conf)out depth_mode conf_modefmapress F/home/cameronsmith/repos/controll3r/dust3r/dust3r/heads/postprocess.py postprocessr sb ;;q!Q "D _T!Q1Q3,%7jI JC$T!Q1*%5IFF Jcb|\}}}|td k(xr|tdk(}|sJ|dk(r|r|S|j||S|jdd}||jdz }|d k(r||jzS|d k(r|t j |zSt d |) z7 extract 3D points from prediction head output inflinear)minmaxT)dimkeepdimg:0yE>)rsquareexp bad mode=)floatclipnormrtorchexpm1 ValueError)xyzrvminvmax no_boundsds rr r sD$%,&BTU5\-AI  x JxxDdx++ R&A 4 C xQXXZ u}U[[^## zD7^ $$rc|\}}}|dk(r&||jj||z zS|dk(r||z tj|z|zSt d|)z8 extract confidence from prediction head output r)rsigmoidr )rr"r$r-r&)xrr(r)s rr r 1smD$ u}aeeglltDyl111 yt u}}Q//$66 zD7^ $$r)r$rr r rrr0s  %6 %r