diagram-type-classifier

BuildwellAI Diagram-Type Classifier

Classifies a UK building-regulations / construction figure into its diagram type, to route a ComfyUI restyle workflow to the class-appropriate prompt/ControlNet.

Classes

• COMPOUND
• CROSS_SECTION
• ISOMETRIC_UPDATE
• SCHEMATIC_FLOW
• SITE_PLAN
• TABLE_DATA

How it was built

• Labels: distilled from a Qwen3.5-VL teacher over 2315 diagram crops (classifications_pool.json).
• Encoder: DINOv3 ViT-L/16 (frozen features for the probe; last blocks fine-tuned for ft_model.pt).
• Finding: encoder size was not the bottleneck (DINOv2-base == DINOv3-large at 234 samples); data scale (234→2315) and fine-tuning to the line-drawing domain were.

Results (5-fold CV / held-out val)

val_acc=0.790 val_bal=0.818

                  precision    recall  f1-score   support

        COMPOUND      0.727     0.762     0.744        84
   CROSS_SECTION      0.855     0.707     0.774       133
ISOMETRIC_UPDATE      0.703     0.867     0.776        30
  SCHEMATIC_FLOW      0.562     0.750     0.643        12
       SITE_PLAN      0.667     0.870     0.755        23
      TABLE_DATA      0.935     0.951     0.943        61

        accuracy                          0.790       343
       macro avg      0.742     0.818     0.772       343
    weighted avg      0.802     0.790     0.791       343

Files

• ft_model.pt — fine-tuned DINOv3 + linear head (state_dict, classes, encoder id).
• class_probe.joblib — frozen-feature logistic-regression probe (lighter, sklearn).
• report.txt — per-class validation report.

Use in ComfyUI

Load DINOv3 + this head in a custom node that maps IMAGE -> (class, confidence) and drives Switch nodes.