Improve model card: pipeline tag, library name, and usage example

README.md

@@ -1,11 +1,12 @@
 ---
+language:
+- en
+license: cc-by-nc-4.0
+pipeline_tag: image-to-image
+library_name: pytorch
 tags:
 - model_hub_mixin
 - pytorch_model_hub_mixin
-license: cc-by-nc-4.0
-language:
-- en
-pipeline_tag: other
 ---
 
 <div align="center">
@@ -25,7 +26,7 @@ pipeline_tag: other
 
 ## Overview
 
-UFM(UniFlowMatch) is a simple, end-to-end trained transformer model that directly regresses pixel displacement image that applies concurrently to both optical flow and wide-baseline matching tasks. 
+UFM(UniFlowMatch) is a simple, end-to-end trained transformer model that directly regresses pixel displacement image that applies concurrently to both optical flow and wide-baseline matching tasks.
 
 This model space contains the refine model.
 
@@ -33,6 +34,45 @@ This model space contains the refine model.
 
 Check out our [Github Repo](https://github.com/UniFlowMatch/UFM) and the hugging face [demo](https://huggingface.co/spaces/infinity1096/UFM).
 
+### Python API
+
+```python
+import cv2
+import torch
+
+# Load the base model (for general use)
+from uniflowmatch.models.ufm import UniFlowMatchConfidence
+model = UniFlowMatchConfidence.from_pretrained("infinity1096/UFM-Base")
+
+# Or load the refinement model (for higher accuracy)
+from uniflowmatch.models.ufm import UniFlowMatchClassificationRefinement
+model = UniFlowMatchClassificationRefinement.from_pretrained("infinity1096/UFM-Refine")
+
+# Set the model to evaluation mode
+model.eval()
+
+# Load images using cv2 or PIL
+source_image = cv2.imread("path/to/source.jpg")
+target_image = cv2.imread("path/to/target.jpg")
+source_rgb = cv2.cvtColor(source_image, cv2.COLOR_BGR2RGB)  # Convert to RGB
+target_rgb = cv2.cvtColor(target_image, cv2.COLOR_BGR2RGB)  # Convert to RGB
+
+# Convert to torch tensors (uint8 or float32)
+# Forward call takes care of normalizing uint8 images appropriate to the UFM model
+source_image = torch.from_numpy(source_rgb)  # Shape: (H, W, 3)
+target_image = torch.from_numpy(target_rgb)  # Shape: (H, W, 3)
+
+# Predict correspondences
+with torch.no_grad():
+    result = model.predict_correspondences_batched(
+        source_image=source_image,
+        target_image=target_image,
+    )
+
+    flow = result.flow.flow_output[0].cpu().numpy()
+    covisibility = result.covisibility.mask[0].cpu().numpy()
+```
+
 
 ## Citation
 If you find our repository useful, please consider giving it a star ⭐ and citing our paper in your work: