Update README.md

README.md

@@ -4,8 +4,11 @@ language:
 license: bsd-3-clause
 tags:
 - arxiv:2506.23151
+- optical-flow-estimation
 pipeline_tag: image-to-image
 library_name: pytorch
+base_model:
+- egorchistov/MEMFOF-Tartan-T
 ---
 
 # MEMFOF-Tartan-T-TSKH
@@ -33,5 +36,16 @@ model = MEMFOF.from_pretrained("egorchistov/MEMFOF-Tartan-T-TSKH").eval().to(dev
 
 with torch.inference_mode():
     example_input = torch.randint(0, 256, [1, 3, 3, 1080, 1920], device=device)  # [B=1, T=3, C=3, H=1080, W=1920]
-    backward_flow, forward_flow = model(example_input).unbind(dim=1)  # [B=1, C=2, H=1080, W=1920]
+    backward_flow, forward_flow = model(example_input)["flow"][-1].unbind(dim=1)  # [B=1, C=2, H=1080, W=1920]
+```
+
+## 📚 Citation
+
+```
+@article{bargatin2025memfof,
+  title={MEMFOF: High-Resolution Training for Memory-Efficient Multi-Frame Optical Flow Estimation},
+  author={Bargatin, Vladislav and Chistov, Egor and Yakovenko, Alexander and Vatolin, Dmitriy},
+  journal={arXiv preprint arXiv:2506.23151},
+  year={2025}
+}
 ```