Skip to content

README.md

Latest commit

 

History

History
218 lines (161 loc) · 8.17 KB

README.md

File metadata and controls

218 lines (161 loc) · 8.17 KB
Raw

English | 中文

IDAMA — Official Code

Implementation for Intermediate Domain Alignment and Morphology Analogy for Patent-Product Image Retrieval (NeurIPS 2025).
Paper / project README · docs/index.html
Dataset & weights Hugging Face — haifan-gong/IDAMA

This folder implements IDM (edge-domain mapping via UAED) and the retrieval stack used with MAF at inference time: preprocessing → feature extraction → similarity & Top-K evaluation → optional MAE/Swin pretraining.

Download PPIRD and checkpoints into IDAMA/data and IDAMA/model (not in git). Set IDAMA_ROOT to the repo root.

Directory layout

code/
├── common/                 # Path constants (paths.py)
├── preprocessing/          # Data & edge preprocessing
│   ├── build_index_and_paths.py    # Path lists & per-item image counts
│   ├── pack_unsupervised_train.py  # Flat UAED train set → sharded tar + manifest
│   ├── pack_sharded_archive.py     # Generic tree → sharded tar (e.g. ImageNet)
│   ├── pack_unsupervised_train.sh
│   ├── pack_features.sh
│   ├── edge_feature/       # UAED / MuGE edge detection
│   │   ├── run_uaed_edge_inference.sh
│   │   └── uaed/           # Train / infer / eval scripts
│   └── legacy/             # Legacy filtering & retrieval helpers
├── feature_extraction/     # Multi-backbone feature extraction
│   ├── extract_features.py           # Unified entry (delegates to legacy)
│   ├── run_feature_extraction.sh     # Extract product + patent features
│   ├── backbones/          # Swin / EVA / EVA02 / iBOT, etc.
│   └── legacy/             # Original extract_feature.py & ResNet variants
├── inference/              # Retrieval inference & metrics
│   ├── compute_similarity.py         # Product–patent similarity matrix
│   ├── rank_and_evaluate.py          # Top-K ranking & evaluation
│   ├── run_inference.sh              # Similarity + eval in one script
│   └── legacy/             # Historical similarity / ranking code
└── pretrain/               # Self-supervised / supervised pretraining
    ├── run_pretrain.sh     # MAE / Swin launcher
    ├── mae/                # Masked Autoencoder
    ├── swin/               # Swin Transformer
    └── ibot/               # iBOT utilities

Requirements

  • Python 3.8+
  • PyTorch + torchvision (CUDA recommended)
  • Common packages: timm, opencv-python, pandas, tqdm, Pillow, numpy

Some submodules need optional deps (e.g. MAE deepspeed, Swin apex); see per-folder READMEs.

export IDAMA_ROOT=/path/to/IDAMA   # default: /data2/gonghaifan/IDAMA
cd "${IDAMA_ROOT}"

Pipeline overview

flowchart LR
  A[Raw images] --> B[UAED edge maps]
  B --> C[Optional: sharded tar packs]
  C --> D[MAE / Swin pretrain]
  B --> E[Multi-backbone features]
  E --> F[Similarity matrix]
  F --> G[Top-K retrieval metrics]
Loading
Stage Module Description
Edge maps preprocessing/edge_feature RGB → edge-domain images for training & retrieval
Packing pack_unsupervised_train.py / pack_sharded_archive.py Large-scale sharded tar + manifest.jsonl
Pretrain pretrain/ Backbone training on edge domain or ImageNet
Features feature_extraction/ Embeddings for test product / patent sets
Inference inference/ Similarity & retrieval metrics

Common commands

Run from the IDAMA repo root (not inside code/). Override paths with IDAMA_ROOT.

1. Build path lists and counts

Regenerate metadata when test image folders change:

python code/preprocessing/build_index_and_paths.py \
  --kind patent \
  --image-dir data/test/new_edge_image/patent \
  --paths-out data/metadata/patent_img_paths.txt \
  --counts-out data/metadata/patent_all.csv

Use --kind product for nested product directory layouts.

2. UAED edge inference

bash code/preprocessing/edge_feature/run_uaed_edge_inference.sh \
  /path/to/raw_images \
  /path/to/edge_output \
  0.5

Default checkpoint: ${IDAMA_ROOT}/model/edge_feature/epoch-19-checkpoint.pth. See preprocessing/edge_feature/README.md.

3. Pack unsupervised training data

Flat UAED directory (patent + product):

bash code/preprocessing/pack_unsupervised_train.sh dry-run   # preview shards
bash code/preprocessing/pack_unsupervised_train.sh run       # write to data/unsupervised_train/

Arbitrary directory tree (e.g. ImageNet):

python code/preprocessing/pack_sharded_archive.py \
  --src /path/to/imagenet/train \
  --out data/unsupervised_train/imagenet1k-edge \
  --num-shards 20 --workers 16 --file-type image

4. Feature extraction

Backbones: r18, r50, r101, swin_t, swin_b, mae, ibot, eva02, clip2, plus _ft fine-tuned variants.

CUDA_VISIBLE_DEVICES=0 bash code/feature_extraction/run_feature_extraction.sh eva02

Output: data/features/new_data_feature/<model_name>/.

5. Similarity & retrieval evaluation

bash code/inference/run_inference.sh eva02

Optional env: POOLING=avgmax|standard, BATCH_SIZE, LOAD_BATCH_NUM.
Writes avgmax_sim.pkl (or avg_sim.pkl) and retrieval/ under the feature directory.

6. Pretraining (MAE / Swin)

cd code/pretrain

# MAE pretrain on edge-domain images
NPROC_PER_NODE=8 bash run_pretrain.sh mae-pretrain \
  "${IDAMA_ROOT}/data/unlabeled_train/data_edge/goods_edge_0.5" \
  "${IDAMA_ROOT}/output/pretrain_mae"

# Swin-Base fine-tune on ImageNet
CUDA_VISIBLE_DEVICES=0,1,2,3 bash run_pretrain.sh swin-train base \
  /path/to/imagenet \
  "${IDAMA_ROOT}/output/swin_base" \
  --resume "${IDAMA_ROOT}/model/swin/pretrained/swin_base_patch4_window7_224.pth"

See run_pretrain.sh for full usage.

Path configuration

common/paths.py defines defaults (edit or remount via IDAMA_ROOT):

Variable Default
METADATA_DIR data/metadata/
FEATURE_DIR data/features/new_data_feature/
PRODUCT_IMAGE_PATHS metadata/product_img_paths.txt
PATENT_IMAGE_PATHS metadata/patent_img_paths.txt

Module reference

preprocessing

  • build_index_and_paths.py: Scan images; write path txt and Index,img_num CSV.
  • pack_unsupervised_train.py: Classify patent/product; parallel tar shards; manifest.jsonl + summary.json.
  • pack_sharded_archive.py: Generic sharded packing with relative paths (ImageNet-style trees).
  • edge_feature/: UAED training & batch inference; Python entry: edge_infer.py.

feature_extraction

  • extract_features.py: Sets sys.path and runs legacy/extract_feature.py.
  • run_feature_extraction.sh: Picks input size, batch, checkpoint; runs product then patent.
  • backbones/: ViT / Swin implementations; legacy/: ResNet and older scripts.

inference

  • compute_similarity.py: --pooling avgmax|standard selects legacy implementation.
  • rank_and_evaluate.py: Top-K ranking and metrics from similarity pickle.
  • run_inference.sh: Runs both steps.

pretrain

  • mae/: MAE pretraining and classification fine-tuning.
  • swin/: Swin classification and SimMIM.
  • ibot/: Weight conversion utilities.

Data layout (under IDAMA_ROOT)

Stage Typical location
Unsupervised tar data/unsupervised_train/ (patent/, product/, imagenet1k-edge/)
Edge training images data/unlabeled_train/data_edge/
Test edge images data/test/new_edge_image/
Feature vectors data/features/new_data_feature/
Retrieval index data/metadata/*.npy, *.csv

See the repo root README and local data/unsupervised_train/README.md (not in git) for dataset notes.

License & attribution

UAED, MAE, Swin Transformer, and other third-party code retain their upstream licenses. This repository is organized for research and reproduction.