search

YangtaoWANG95 / TokenCut

(CVPR 2022) Pytorch implementation of "Self-supervised transformers for unsupervised object discovery using normalized cut"
MIT License
304 stars 34 forks source link

Is it real-time capable of doing segmentation on mobile devices? #9

Closed roimulia2 closed 2 years ago

roimulia2 commented 2 years ago

Hey there! Love the paper.

Is it real-time capable of doing segmentation on mobile devices?

Thanks!

YangtaoWANG95 commented 2 years ago

Hi,

We haven't tried to test on mobile devices. As mentioned on the paper, for 480p images, it takes about 0.32s. If you want to make the inference faster. You can try to use torch Eigen decomposition function https://pytorch.org/docs/stable/generated/torch.lobpcg.html to replace the spicy Eigen decomposition https://github.com/YangtaoWANG95/TokenCut/blob/0d86de5ab91b846e730e5e4cfee510544e53407b/object_discovery.py#L40

Here is a pseudocode: A = A > tau A = A.float() A = A + eps d_i = torch.sum(A, dim=2) D = torch.diag_embed(d_i, dim1=1) X = (D-A) / (D + eps) eigval, eigvec = torch.lobpcg(A=D-A, B=D, k=2, largest=False) second_smallest_vec = eigvec[0,:,1].cpu().numpy()

Best, Yangtao