Add metadata and improve model card

README.md

@@ -1,3 +1,8 @@
+---
+license: apache-2.0
+pipeline_tag: image-segmentation
+---
+
 <div align="center">
 
 <h1>VGGT-Segmentor: Geometry-Enhanced Cross-View Segmentation</h1>
@@ -47,6 +52,7 @@
 
 </div>
 
+This repository contains the official weights for the paper [VGGT-Segmentor: Geometry-Enhanced Cross-View Segmentation](https://arxiv.org/abs/2604.13596).
 
 ## 📢 News
 
@@ -122,305 +128,13 @@ hf download zbbhhh/VGGT-S main_exp.pth --local-dir official_ckpts # Download che
 
 ### Modify VGGT Track Head
 
-To extract VGGT encoder feature maps for downstream correspondence learning, please modify the `forward` function in:
-
-```text
-third_party/vggt_main/vggt/heads/track_head.py
-```
-
-Replace the original `forward` function (around Line 72) with the following implementation:
-
-```python
-def forward(self, aggregated_tokens_list, images, patch_start_idx, query_points=None, iters=None, feat=False):
-    """
-    Forward pass of the TrackHead.
-
-    Args:
-        aggregated_tokens_list (list): List of aggregated tokens from the backbone.
-        images (torch.Tensor): Input images of shape (B, S, C, H, W), where:
-                                B = batch size,
-                                S = sequence length.
-        patch_start_idx (int): Starting index for patch tokens.
-        query_points (torch.Tensor, optional): Initial query points to track.
-                                                If None, points are initialized automatically.
-        iters (int, optional): Number of refinement iterations.
-                                If None, uses self.iters.
-        feat (bool, optional): Whether to additionally return feature maps.
-
-    Returns:
-        tuple:
-            - coord_preds (torch.Tensor): Predicted coordinates.
-            - vis_scores (torch.Tensor): Visibility scores.
-            - conf_scores (torch.Tensor): Confidence scores.
-            - feature_maps (torch.Tensor, optional): Extracted VGGT feature maps.
-    """
-
-    B, S, _, H, W = images.shape
-
-    # Extract feature maps from VGGT tokens
-    # Shape: (B, S, C, H//2, W//2)
-    feature_maps = self.feature_extractor(
-        aggregated_tokens_list,
-        images,
-        patch_start_idx
-    )
-
-    # Use default iterations if not specified
-    if iters is None:
-        iters = self.iters
-
-    # Tracking
-    coord_preds, vis_scores, conf_scores = self.tracker(
-        query_points=query_points,
-        fmaps=feature_maps,
-        iters=iters
-    )
-
-    if feat:
-        return coord_preds, vis_scores, conf_scores, feature_maps
-
-    return coord_preds, vis_scores, conf_scores
-```
-
-#### Optional
-
-We observe that VGGT recomputes the 2D positional embeddings during every forward pass, even when the input image resolution remains unchanged. To accelerate both training and inference, we adapt the implementation by precomputing the 2D positional embeddings offline and loading them directly at runtime, thereby avoiding redundant computations.
-
-1. Generate the 2D positional embeddings using `gen_pos_embed.py`:
-
-```bash
-python gen_pos_embed.py --img_H 518 --img_W 518
-```
-
-We also provide precomputed positional embeddings for various commonly used image resolutions on [Hugging Face](https://huggingface.co/zbbhhh/VGGT-S).
-
-2. Modify the VGGT tracking head to use the precomputed embeddings.
-
-Open:
-
-```text
-third_party/vggt_main/vggt/heads/track_modules/base_track_predictor.py
-```
-
-Add the following line at the end of the `__init__` method in the `BaseTrackerPredictor` class, around [base_track_predictor.py#L81](https://github.com/facebookresearch/vggt/blob/main/vggt/heads/track_modules/base_track_predictor.py#L81):
-
-```python
-self.pos_embed = None
-```
-
-This attribute is used to cache the loaded positional embeddings.
-
-Then, replace the following line in [base_track_predictor.py#L149](https://github.com/facebookresearch/vggt/blob/main/vggt/heads/track_modules/base_track_predictor.py#L149):
-
-```python
-pos_embed = get_2d_sincos_pos_embed(self.transformer_dim, grid_size=(HH, WW)).to(query_points.device)
-```
-
-with:
-
-```python
-pos_emb_pt = (
-    f"vggt_main/pos_embed/"
-    f"pos_embed_{self.transformer_dim}_{HH}_{WW}.pt"
-)
-
-assert os.path.exists(pos_emb_pt), (
-    f"[BUG] Positional embedding file not found: {pos_emb_pt}"
-)
-
-if self.pos_embed is None:
-    self.pos_embed = torch.load(pos_emb_pt).to(query_points.device)
-
-pos_embed = self.pos_embed
-```
+To extract VGGT encoder feature maps for downstream correspondence learning, please modify the `forward` function in `third_party/vggt_main/vggt/heads/track_head.py`. Replace the original `forward` function with the implementation provided in the [official GitHub repository](https://github.com/buaa-colalab/VGGT-S).
 
 ---
 
 ### Dataset Preparation
 
-#### Download Ego-Exo4D
-
-Please follow the official [SegSwap Ego-Exo4D correspondence pipeline](https://github.com/EGO4D/ego-exo4d-relation/tree/main/correspondence/SegSwap) to:
-
-1. Download the Ego-Exo4D dataset
-2. Extract video frames into images
-3. Generate correspondence pairs
-
-After preprocessing, the dataset should be organized as:
-
-```text
-data_root/
-├── take_id_01/
-│   ├── ego_cam/
-│   │   ├── 0.jpg
-│   │   ├── ...
-│   │   └── n.jpg
-│   ├── exo_cam/
-│   │   ├── 0.jpg
-│   │   ├── ...
-│   │   └── n.jpg
-│   └── annotation.json
-├── ...
-├── take_id_n/
-└── split.json
-```
-
-Then run the official SegSwap pair-generation script:
-
-[create_pairs.py](https://github.com/EGO4D/ego-exo4d-relation/blob/main/correspondence/SegSwap/data/create_pairs.py)
-
-This will generate:
-
-```text
-train_egoexo_pairs.json
-train_exoego_pairs.json
-val_egoexo_pairs.json
-val_exoego_pairs.json
-```
-
----
-
-#### Configure Dataset Path
-
-If the Ego-Exo4D dataset has already been downloaded, modify:
-
-```text
-src/dataloader.py
-```
-
-and set:
-
-```python
-EGOEXO4D_ROOT = "your/data/path"
-```
-
----
-
-### Data Preprocessing
-> We provide all intermediate result files generated from the following preprocessing steps. Therefore, if you have already downloaded the Ego-Exo4D dataset, you can directly use the provided files without rerunning the preprocessing pipeline. Nevertheless, we still release the complete data preprocessing code to support further community research and help readers better understand our pipeline and methodology.
-
-#### Generate Scene JSON
-This step reorganizes the dataset structure and converts the pair_json to the following hierarchical JSON format:
-
-```text
-scene
-└── [obj_info]
-    ├── ego_info
-    │   ├── ego_rgb
-    │   └── ego_mask
-    └── exo_info
-        ├── exo_rfb
-        └── exo_mask
-```
-
-where each `scene` contains multiple object annotations (`obj_info`), and each object stores the corresponding ego-view and exo-view information, including RGB frames (`*_rgb`) and segmentation masks (`*_mask`). Meanwhile, only objects that simultaneously appear in both ego and exo views are retained.
-
-
-```bash
-cd data
-python gen_scenes.py
-
-# Or, downloading from huggingface
-huggingface-cli download zbbhhh/VGGT-S train_scenes.json --local-dir data
-```
-
-Outputs:
-
-```text
-train_scenes.json
-val_scenes.json
-```
-
----
-
-#### Generate Single-Object JSON
-This step converts `train_scenes.json` and `val_scenes.json` into a sample-based format, where each individual object instance is treated as a training sample for convenient dataloader indexing and loading.
-
-Each sample follows the format:
-
-```text
-scene_id//ego_cam//exo_cam//object//frame_id
-```
-
-For example:
-
-```text
-c9e0bbae-4092-4ab1-95cd-16b905709a0e//aria01_214-1//gp01//white spatula_0//002160
-```
-
-which represents:
-
-* `scene_id`: `c9e0bbae-4092-4ab1-95cd-16b905709a0e`
-* `ego_cam`: `aria01_214-1`
-* `exo_cam`: `gp01`
-* `object`: `white spatula_0`
-* `frame_id`: `002160`
-
-```bash
-python gen_one_object.py
-
-# Or, downloading from huggingface
-huggingface-cli download zbbhhh/VGGT-S train_obj.json --local-dir data
-```
-
-Outputs:
-
-```text
-train_obj.json
-val_obj.json
-```
-
----
-
-#### Generate VGGT Projected Points
-This step precomputes the VGGT projected points as the first-stage mapping points and saves them offline. By caching these mappings in advance, the training and inference stages can directly perform the second-stage mapping without rerunning the initial VGGT projection process, significantly reducing computation time.
-
-Optionally, the first-stage mapping can also be performed online if needed.
-
-##### Ego → Exo
-
-```bash
-python extract_point_g2x.py
-```
-
-Outputs:
-
-```text
-train_obj_wp_g2x.json
-val_obj_wp_g2x.json
-```
-
-##### Exo → Ego
-
-```bash
-python extract_point_x2g.py
-```
-
-Outputs:
-
-```text
-train_obj_wp_x2g.json
-val_obj_wp_x2g.json
-```
-
----
-
-#### Merge Point Annotations
-This step merges the projected points generated in the previous step and produces the final data annotation JSON file used for training and inference.
-
-```bash
-python merge_point.py
-
-# Or, downloading from huggingface
-huggingface-cli download zbbhhh/VGGT-S train_obj_wp.json --local-dir data
-```
-
-Final outputs:
-
-```text
-train_obj_wp.json
-val_obj_wp.json
-```
+Please follow the official [SegSwap Ego-Exo4D correspondence pipeline](https://github.com/EGO4D/ego-exo4d-relation/tree/main/correspondence/SegSwap) to download and preprocess the Ego-Exo4D dataset.
 
 ---
 
@@ -443,10 +157,8 @@ python -m torch.distributed.run \
     --config hybrid.yaml
 ```
 
-> Training and validation only require modifying the parameters in `hybrid.yaml`, where all configuration parameters are carefully explained.
 ---
 
-
 ## 📝 Citation
 
 If you find this work useful, please consider citing our paper:
@@ -466,11 +178,4 @@ This project is licensed under the Apache-2.0 License. See [LICENSE](LICENSE.txt
 
 ## 🙏 Acknowledgement
 
-This project is built upon several excellent open-source projects:
-
-* [VGGT](https://github.com/facebookresearch/vggt)
-* [EgoExo4d](https://github.com/EGO4D/ego-exo4d-relation)
-* [SAM2](https://github.com/facebookresearch/sam2)
-
-We thank the authors for releasing their code and models to the community.
-
+This project is built upon several excellent open-source projects: [VGGT](https://github.com/facebookresearch/vggt), [EgoExo4d](https://github.com/EGO4D/ego-exo4d-relation), and [SAM2](https://github.com/facebookresearch/sam2).