MatteoCargnelutti's picture
Update README.md
4781ba4 verified
|
Raw
History Blame Contribute Delete
13.3 kB
metadata
tags:
  - efficient-net
  - image-classification
  - document-analysis
  - visual-elements
  - cultural-heritage
  - digitized-books
datasets:
  - institutional/institutional-books-hl-visual-elements
pipeline_tag: image-classification
model-index:
  - name: institutional-books-visual-elements-orientation
    results:
      - task:
          type: image-classification
          name: Visual Element Orientation Classification
        metrics:
          - type: accuracy
            name: Overall accuracy
            value: 0.9134
            dataset:
              type: institutional/institutional-books-hl-visual-elements
              name: Visual elements crops (orientation test split)
          - type: accuracy
            name: upright accuracy
            value: 0.9178
          - type: accuracy
            name: rotate_90_clockwise accuracy
            value: 0.8976
          - type: accuracy
            name: rotate_180 accuracy
            value: 0.9191
          - type: accuracy
            name: rotate_90_counterclockwise accuracy
            value: 0.9191
license: apache-2.0

📚 Institutional Books — Visual Elements Orientation Model (EfficientNetV2-M)

A 4‑class image classification model that predicts the rotation correction needed to restore visual elements from digitized book page scans to upright orientation. This model operates on cropped regions (e.g., images, diagrams, ornaments) and is intended as a post-processing stage after visual-element detection.

More information:

See also:

The Institutional Data Initiative at Harvard Law School Library works with knowledge institutions—from libraries and museums to cultural groups and government agencies—to refine and publish their collections as data. Reach out to collaborate on your collections.


Outline


Model Description

  • Architecture: EfficientNetV2-M
  • Parameters: 52,863,480
  • Classes (corrections): 4
  • Input resolution (train/inference): 480×480 px crops (from 512×512 resize)
  • Framework: PyTorch

The model predicts the inverse rotation required to correct each crop back to upright. Input crops may be in any of four orientations; the model outputs one of four rotation labels.

Classes

The labels are “correction actions” to make the image upright:

Class index Label Description
0 upright No rotation needed (already upright)
1 rotate_90_clockwise Rotate 90° clockwise to correct
2 rotate_180 Rotate 180° to correct
3 rotate_90_counterclockwise Rotate 90° counter-clockwise to correct

Classes are mutually exclusive.

Intended Use

This model classifies the orientation of already-detected visual elements from digitized book pages.

Primary use cases:

  • Correcting rotations of cropped visual elements in digitization pipelines
  • Normalizing orientation before downstream tasks (captioning, OCR on musical scores, layout analysis)
  • Quality control on large-scale digitized collections (flagging mis-rotate elements)

Out of scope:

  • General full-page orientation detection (expects tight crops of a single element)
  • Classification of semantic content (e.g., “image vs. music”)
  • Arbitrary-angle rotation beyond multiples of 90°

Training

Dataset

Source data consists of 7,904 manually curated crops of visual elements from the Institutional Books collection.

Original (pre-synthetic) label distribution (estimated, by source orientation):

  • upright: 93.0%
  • rotate_90_clockwise: 5.6%
  • rotate_90_counterclockwise: 1.3%
  • rotate_180: 0.03%

To avoid this extreme imbalance and to directly learn the correction operation:

  1. All images were first manually corrected to upright.
  2. Each upright image was synthetically rotated by 0°, 90°, 180°, and 270°.
  3. The target label is the inverse rotation needed to restore the upright orientation.

Resulting synthetic orientation dataset:

  • Total samples (with synthetic rotations): 31,616
  • Train samples: 25,292
  • Val samples: 3,160
  • Test samples: 3,164
  • Split: 0.8 / 0.1 / 0.1 (train / val / test)

Each original crop contributes four synthetic samples (one per orientation), producing a balanced label distribution across the four classes in the synthetic set.

Training Configuration

Parameter Value
Backbone EfficientNetV2-M
Classifier head Dropout(p=0.3) → Linear(1280, 4)
Image size (train) Resize(512×512) → RandomCrop(480×480)
Image size (val/test) Resize(512×512) → CenterCrop(480×480)
Batch size 32
Max epochs 20
Optimizer / LR Not specified (standard schedule)
Normalization ImageNet mean/std
Hardware Single NVIDIA GH200 GPU
Total training time 58 min 13 sec (20 epochs)
Avg. per epoch ~2 min 55 sec
Train samples/epoch 25,292 (~791 steps/epoch)
Throughput ~145 images/sec

Data Augmentation (Train Only)

Preprocessing:

  • Resize(512×512)
  • RandomCrop(480×480)
  • ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1)

Stochastic augmentations:

Augmentation Implementation Probability
Random auto-contrast RandomAutocontrast p = 0.3
Random invert RandomInvert p = 0.15
Random grayscale RandomGrayscale p = 0.2
Gaussian blur GaussianBlur(k=5, σ=0.1–2.0) p = 1.0
Random erasing RandomErasing(scale=0.02–0.15) p = 0.2

Validation/Test preprocessing:

  • Resize(512×512)
  • CenterCrop(480×480)
  • Normalize (ImageNet stats)

Evaluation

Per-Epoch Training Summary

On the validation set, accuracy increases steadily and plateaus around 90%, while validation loss begins to rise after approximately epoch 10, indicating moderate overfitting. Heavy augmentations successfully limit overfitting enough that the held-out test set slightly outperforms validation.

Final epoch (20):

  • Train accuracy: 99.81%
  • Val accuracy: 90.35%
  • Train loss: 0.0058
  • Val loss: 0.4753

Test Set Performance

On the held-out test set (3,164 samples; 791 per class):

  • Overall accuracy: 91.34% (2,890 / 3,164)

Per-class accuracy:

Class Accuracy Correct / Total
upright 91.78% 726 / 791
rotate_90_clockwise 89.76% 710 / 791
rotate_180 91.91% 727 / 791
rotate_90_counterclockwise 91.91% 727 / 791

Additional notes:

  • Misclassifications: 274 of the 3,164 test samples are misclassified (2,890 correct → 91.34% overall accuracy). These break down per class as 65 (upright), 81 (rotate_90_clockwise), 64 (rotate_180), and 64 (rotate_90_counterclockwise).
  • The 90° clockwise class is the most challenging, but still achieves close to 90% accuracy.

Inference Configuration

Typical inference settings:

Parameter Value
Image size 512×512 resize → 480×480 center crop
Batch size 32 (tune for available GPU memory)
Normalization ImageNet mean/std
Output 4-way softmax over orientation labels

The top-1 prediction corresponds to the rotation to apply to make the crop upright.

Usage

PyTorch Example

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import transforms
from huggingface_hub import hf_hub_download
from PIL import Image

# Download weights (the repo ships a state_dict at weights/weights.pth)
model_path = hf_hub_download(
    repo_id="institutional/institutional-books-visual-elements-orientation",
    filename="weights/weights.pth",
)

# Build the architecture and load the state_dict
model = models.efficientnet_v2_m(weights=None)
num_features = model.classifier[1].in_features
model.classifier = nn.Sequential(
    nn.Dropout(p=0.3, inplace=True),
    nn.Linear(num_features, 4),
)
state_dict = torch.load(model_path, map_location="cuda", weights_only=True)
model.load_state_dict(state_dict)
model.to("cuda")
model.eval()

# Preprocessing: match validation/test pipeline
preprocess = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.CenterCrop(480),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],  # ImageNet
        std=[0.229, 0.224, 0.225],
    ),
])

idx_to_label = {
    0: "upright",
    1: "rotate_90_clockwise",
    2: "rotate_180",
    3: "rotate_90_counterclockwise",
}

# A high confidence threshold (0.99) is applied: predictions below it
# default to "upright" to minimize false corrections.
CONFIDENCE_THRESHOLD = 0.99

def predict_orientation(path):
    img = Image.open(path).convert("RGB")
    x = preprocess(img).unsqueeze(0).to("cuda")
    with torch.no_grad():
        logits = model(x)
        probs = torch.softmax(logits, dim=1)[0]
    top1 = int(torch.argmax(probs))
    conf = float(probs[top1])
    label = idx_to_label[top1] if conf >= CONFIDENCE_THRESHOLD else "upright"
    return label, conf, probs.cpu().tolist()

label, conf, all_probs = predict_orientation("crop.jpg")
print(f"Predicted correction: {label}, confidence: {conf:.3f}")

Applying Corrections

from PIL import Image

def apply_correction(img, label):
    if label == "upright":
        return img
    elif label == "rotate_90_clockwise":
        return img.rotate(-90, expand=True)
    elif label == "rotate_180":
        return img.rotate(180, expand=True)
    elif label == "rotate_90_counterclockwise":
        return img.rotate(90, expand=True)
    else:
        raise ValueError(f"Unknown label: {label}")

Limitations

  • Trained specifically on crops from the Institutional Books collection. Performance may degrade on:
    • Non-book imagery
    • Heavily stylized or abstract content
    • Very low-resolution or heavily compressed scans
  • Supports only multiples of 90° rotations; does not handle slight skews or arbitrary angle rotations.
  • Expected to work best when:
    • Crops contain a clear visual object/structure
    • Background is not overwhelmingly dominant
  • Model assumes images are RGB; grayscale images are internally handled via standard preprocessing but not natively optimized for non-RGB channels.

Citation

@misc{mendez2026institutionalbooksvisual,
      title={Institutional Books - Visual Elements: An open-source pipeline for extracting, classifying, deduplicating, and captioning visual elements from digital book collections}, 
      author={Jimmy Mendez and Matteo Cargnelutti and David Lowry-Duda and Catherine Brobston and Salwa Ismail and Greg Leppert and Amanda Watson and Jonathan Zittrain},
      year={2026},
      eprint={2608.18957},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2608.18957}, 
}