BPBReID (Body Part-based Re-Identification)

repository·main·Indexed 42 days ago

https://github.com/vlsomers/bpbreid

A body part-based representation learning framework for person re-identification designed for robustness against occlusions. Built on top of Torchreid, it uses an attention mechanism to pool spatial feature maps into multiple body part-based embeddings and employs visibility-aware similarity scoring. Key components include the ImagePartBasedEngine, GiLt loss, and BodyPartAttentionLoss.

Tokens
15.3K
Snippets
33
Records
67
Agent score
51%

What's inside BPBReID

  1. Overview of BPBReID

    main

    BPBReID (Body Part-based Re-Identification) is a method for person re-identification designed to be robust against occlusions. Unlike standard global ReID methods that output a single feature per sample, BPBReID outputs multiple features—one for each body part.

    Key Concepts:

    • Part-based Representation: Uses an attention mechanism to pool global spatial feature maps into $K$ body part-based embeddings.
    • Occlusion Robustness: At test time, similarity scores are computed using the average distance of only the mutually visible part-based embeddings. Visibility scores are derived from attention map activations.
    • Training: Uses pseudo human parsing labels to learn the attention mechanism, but no human parsing labels are required during inference.
  2. Explore the Awesome-ReID collection of research

    main

    The Awesome-ReID document provides a curated list of Person Re-identification (ReID) research papers and their corresponding code implementations. It serves as a resource for developers and researchers to find state-of-the-art methods across various sub-fields, including:

    • Cross-Modality ReID: e.g., RGB-Infrared alignment.
    • Unsupervised/Self-Supervised ReID: e.g., Graph association, clustering, and domain adaptation.
    • Video-based ReID: e.g., Spatial-temporal modeling and attention networks.
    • Attention Mechanisms: e.g., Mixed high-order attention and co-segmentation inspired networks.
    • Domain Adaptation: e.g., Camera style adaptation and exemplar memory.

    Each entry in the collection typically includes a link to the research paper (ArXiv or CVF) and a link to the official or community-maintained code repository.

  3. Use torchreid.data for data management

    main
    The torchreid.data module provides the core infrastructure for handling Re-Identification datasets. It includes components for managing data loading, sampling strategies, image/video transformations, and specific dataset implementations.
  4. Use Image Engines for image-based ReID training

    main

    For training Re-Identification models using static images, use one of the following specialized image engines:

    • torchreid.engine.image.softmax.ImageSoftmaxEngine: Optimized for training with a Softmax loss function.
    • torchreid.engine.image.triplet.ImageTripletEngine: Optimized for training with a Triplet loss function.
  5. Core components of the BPBReID framework

    main

    The BPBReID codebase is built on top of the Torchreid framework and introduces several specialized components for part-based ReID:

    • ImagePartBasedEngine (torchreid/engine/image/part_based_engine.py): Handles training and testing of part-based models. It manages query-gallery distance matrices using multiple features per sample and incorporates visibility scores.
    • GiLt loss (/torchreid/losses/GiLt_loss.py): A configurable loss function that selectively applies ID/triplet loss to both holistic (global) and part-based features.
    • BodyPartAttentionLoss (torchreid/losses/body_part_attention_loss.py): Specifically used to train the attention mechanism.
    • BPBreID (torchreid/models/bpbreid.py): The core model supporting various part-based feature computation methods, including learnable attention, fixed attention heatmaps from external models, and PCB-like horizontal stripes.
    • Logger (torchreid/utils/logging/logger.py): Supports logging to Weights & Biases and other tools.
    • EngineState (torchreid/utils/engine_state.py): Tracks training progress (e.g., current epoch).
    • visualize_query_gallery_rankings.py (torchreid/utils/visualization/visualize_query_gallery_rankings.py): A tool for ranking visualization that displays part heatmaps and local distances for each part.
  6. Use the ranking visualization tool

    main

    The ranking visualization tool allows you to inspect Top-K ranked samples and their corresponding attention maps.

    Configuration:

    • Enable the tool by setting test.visrank to True in your config.
    • Set the number of samples to display using test.visrank_topk.

    Visualization Details:

    • Rows: The first row (blue background) is the query. Subsequent rows (green/red) indicate correct/incorrect matches.
    • Attention Maps: Displays foreground and part-based attention.
      • Borders: Green borders indicate visible parts; red borders indicate unvisible parts.
      • Metrics: The first number under an attention map is the visibility score; the second is the distance to the query embedding.
    • Global Distances: Displayed under images in the first column, typically computed as the average of all other distances weighted by visibility scores.
  7. Use Video Engines for video-based ReID training

    main

    For training Re-Identification models using video sequences, use one of the following specialized video engines:

    • torchreid.engine.video.softmax.VideoSoftmaxEngine: Optimized for training video data with a Softmax loss function.
    • torchreid.engine.video.triplet.VideoTripletEngine: Optimized for training video data with a Triplet loss function.
  8. Quickstart: Train and test a model in Python

    main

    You can use Torchreid's unified API to perform end-to-end training and evaluation. The workflow involves:

    1. Initializing an ImageDataManager to handle datasets.
    2. Building a model using torchreid.models.build_model.
    3. Building an optimizer and lr_scheduler using torchreid.optim.
    4. Initializing an ImageSoftmaxEngine to manage the training loop.
    5. Calling engine.run() to execute training and testing.
    import torchreid
    
    # 1. Load data manager
    datamanager = torchreid.data.ImageDataManager(
        root='reid-data',
        sources='market1501',
        targets='market1501',
        height=256,
        width=128,
        batch_size_train=32,
        batch_size_test=100,
        transforms=['random_flip', 'random_crop']
    )
    
    # 2. Build model, optimizer and lr_scheduler
    model = torchreid.models.build_model(
        name='resnet50',
        num_classes=datamanager.num_train_pids,
        loss='softmax',
        pretrained=True
    )
    model = model.cuda()
    
    optimizer = torchreid.optim.build_optimizer(
        model,
        optim='adam',
        lr=0.0003
    )
    
    scheduler = torchreid.optim.build_lr_scheduler(
        optimizer,
        lr_scheduler='single_step',
        stepsize=20
    )
    
    # 3. Build engine
    engine = torchreid.engine.ImageSoftmaxEngine(
        datamanager,
        model,
        optimizer=optimizer,
        scheduler=scheduler,
        label_smooth=True
    )
    
    # 4. Run training and test
    engine.run(
        save_dir='log/resnet50',
        max_epoch=60,
        eval_freq=10,
        test_only=False
    )
  9. Visualize ranking results

    main

    To visualize how the model ranks images, set visrank=True in engine.run(). This requires test_only=True.

    • visrank_topk: Determines how many top images to visualize (default is 10).
    • Output: Results are saved in save_dir/visrank_DATASETNAME. Plots show correct matches in green and incorrect matches in red.
  10. Perform two-stepped transfer learning

    main

    To prevent pretrained layers from being damaged by gradients from randomly initialized layers, you can use a two-stepped strategy in engine.run():

    1. Freeze base layers: Set fixbase_epoch to the number of epochs you want to keep the base layers frozen.
    2. Open specific layers: Use open_layers to specify which layers (e.g., 'classifier') should be trained during the frozen period.

    Note that fixbase_epoch is included in the total max_epoch count.

    engine.run(
        save_dir='log/resnet50',
        max_epoch=60,
        eval_freq=10,
        test_only=False,
        fixbase_epoch=5,
        open_layers='classifier'
    )
  11. Configure and use human parsing labels

    main

    BPBReID requires human parsing labels (masks) for datasets.

    Using Pre-generated Labels

    You can download pre-generated labels for Market-1501, DukeMTMC-reID, Occluded-Duke, Occluded-ReID, and P-DukeMTMC from Google Drive.

    After unzipping, place the masks folder inside your dataset directory following this structure:

    Market-1501-v15.09.15
    ├── bounding_box_test
    ├── bounding_box_train
    ├── masks
    │   └── pifpaf_maskrcnn_filtering
    │       ├── bounding_box_test
    │       ├── bounding_box_train
    │       └── query
    └── query

    Important: Set the data.root configuration key to the parent directory containing all your dataset folders.

    Generating Labels for Custom Datasets

    To generate labels for a new dataset, run the get_labels script. The labels will be saved in a masks folder within your source directory.

    conda activate bpbreid
    python torchreid/scripts/get_labels --source [Dataset Path]