nithish-007/mi_jepa_based_clustering

How nithish-007/mi_jepa_based_clustering works

This repository implements an SKU/product-image clustering workflow: prepare image data, train a representation model, cluster unseen products, and evaluate clustering and retrieval quality. The model combines a ViT backbone, invariance and discrimination heads, and an optional I-JEPA training branch; inference embeds images and offers graph, HDBSCAN, and oracle k-means methods. Data-module implementations and some pipeline details are not shown, so their internal connections are omitted; edges reflect the supplied excerpts, source index, or file names.

Main entry points: network.py Clustering network · module.py Training module · pipeline.py Evaluation pipeline · embedding.py Model embedding

Components 26

Product workflow

Config loader src/config.py
RP2K download scripts/download_rp2k.py
Data preparation scripts/prepare_data.py
Image data src/data
Train command scripts/train.py
Cluster command scripts/cluster.py
Evaluate command scripts/evaluate.py

Model and training

Training runner src/training/trainer.py
Training module src/training/module.py
ViT backbone src/models/backbone.py
Clustering network src/models/network.py
Embedding heads src/models/heads.py
I-JEPA branch src/models/ijepa.py
Training losses src/losses
Checkpoint handling src/training/checkpoints.py

Inference and clustering

Inference runner src/inference/run.py
Model embedding src/inference/embedding.py
Method dispatch src/inference/pipeline.py
Pair scoring src/inference/scoring.py
Clustering methods src/inference/clustering.py
Cluster export src/inference/export.py

Evaluation

Evaluation pipeline src/evaluation/pipeline.py
Metrics src/evaluation/metrics.py
Evaluation plots src/evaluation/plots.py
Model analysis src/evaluation/analysis.py

Other components

Researcher

Connections 40

Researcher → RP2K download: obtains data
Researcher → Data preparation: prepares data
Researcher → Train command: starts training
Researcher → Cluster command: runs clustering
Researcher → Evaluate command: runs evaluation
Data preparation → Config loader: loads settings scripts/prepare_data.py
Data preparation → Image data: prepares images
RP2K download → Image data: downloads images
Train command → Config loader: loads settings scripts/train.py
Train command → Training runner: starts training scripts/train.py
Training runner → Training module: trains module src/training/trainer.py
Training module → Clustering network: forwards images src/training/module.py
Training module → Training losses: computes objectives src/training/module.py
Clustering network → ViT backbone: extracts features src/models/network.py
Clustering network → Embedding heads: projects features src/models/network.py
Clustering network → I-JEPA branch: creates branch src/models/network.py
Training runner → Checkpoint handling: manages checkpoints src/training/trainer.py
Cluster command → Config loader: loads settings scripts/cluster.py
Cluster command → Inference runner: starts inference scripts/cluster.py
Inference runner → Model embedding: loads model src/inference/run.py
Model embedding → ViT backbone: builds baseline src/inference/embedding.py
Model embedding → Clustering network: builds trained model src/inference/embedding.py
Model embedding → Checkpoint handling: resolves checkpoint src/inference/embedding.py
Inference runner → Method dispatch: dispatches methods src/inference/run.py
Method dispatch → Clustering methods: runs baselines src/inference/pipeline.py
Method dispatch → Pair scoring: scores graph pairs src/inference/pipeline.py
Pair scoring → Clustering methods: clusters candidates src/inference/scoring.py
Inference runner → Cluster export: exports results src/inference/run.py
Evaluate command → Config loader: loads settings scripts/evaluate.py
Evaluate command → Evaluation pipeline: starts evaluation scripts/evaluate.py
Evaluation pipeline → Model embedding: embeds products src/evaluation/pipeline.py
Evaluation pipeline → Method dispatch: predicts clusters src/evaluation/pipeline.py
Evaluation pipeline → Metrics: scores predictions src/evaluation/pipeline.py
Evaluation pipeline → Evaluation plots: saves diagnostics src/evaluation/pipeline.py
Training module → Metrics: validates retrieval src/training/module.py
Training module → Pair scoring: fits calibration src/training/module.py
Model analysis → Metrics: uses metrics src/evaluation/analysis.py
Model analysis → Clustering methods: uses clustering src/evaluation/analysis.py
Model analysis → Clustering network: analyzes model src/evaluation/analysis.py
Model analysis → Training losses: analyzes objectives src/evaluation/analysis.py

Drawn by GitDiagram on November 14, 2026 from the repository's file tree, README and source files. Read as Markdown ·