capturemate-image-classifier-v1-5class

CaptureMate Image Classifier v1 5-Class

CLIP-based screenshot image classifier for CaptureMate.

This model classifies screenshot images into one of five CaptureMate categories:

• schedule
• shopping
• place
• memo
• unknown

The model is designed to complement the CaptureMate OCR text classifier by providing visual information when screenshot text alone is ambiguous or incomplete.

Note: This model uses a custom PyTorch classification head on top of a CLIP vision encoder. Therefore, it cannot be loaded directly with AutoModelForImageClassification or the standard Hugging Face image-classification pipeline.

Model Details

Model Description

This model is a CLIP-based image classification model developed for CaptureMate, an iOS screenshot organization and action recommendation app.

It receives a screenshot image and predicts the most relevant CaptureMate category based on visual information.

• Developed by: CaptureMate
• Model type: Image classification
• Base image encoder: openai/clip-vit-base-patch32
• Architecture: CLIP vision encoder with a custom PyTorch classification head
• Framework: PyTorch
• Task: Single-label multi-class image classification
• Input: Screenshot image
• Output: One of five CaptureMate categories
• Number of labels: 5
• License: Not specified

Version

This repository contains v1 of the 5-class CaptureMate screenshot image classifier.

The model predicts the following categories:

• schedule
• shopping
• place
• memo
• unknown

Future versions may improve image-only classification performance as additional screenshot data becomes available.

Labels

Intended Use

Direct Use

This model can be used to classify screenshot images into the five CaptureMate categories.

Unlike the CaptureMate OCR text classifier, this model analyzes the visual information contained in the screenshot.

The model is especially useful when:

• OCR text is short or incomplete
• OCR extraction contains errors
• The screenshot contains useful visual structure
• Text-based classification is uncertain
• Multiple categories have similar textual content but different visual characteristics

Multimodal Use

This model is primarily designed to be used together with the CaptureMate OCR text classifier.

The overall CaptureMate classification pipeline uses both OCR text and screenshot image information:

                    Screenshot
                        |
             +----------+----------+
             |                     |
             v                     v
        OCR Extraction       Image Analysis
             |                     |
             v                     v
      Text Classifier       Image Classifier
             |                     |
             +----------+----------+
                        |
                        v
                  Score Fusion
                        |
                        v
                Category Decision
                        |
                        v
             Action Recommendation

In the current CaptureMate multimodal pipeline, the prediction scores are combined using fixed weighted fusion:

Text weight  = 0.7
Image weight = 0.3

Conceptually:

final_score =
    0.7 × text_score
    +
    0.3 × image_score

The category with the highest fused score is selected as the final category prediction.

The image classifier is therefore intended to provide a secondary visual signal rather than replace the OCR text classifier.

Out-of-Scope Use

This model is not designed to:

• Perform OCR
• Extract text from screenshots
• Classify arbitrary non-screenshot images
• Extract structured fields such as dates, prices, addresses, or product names
• Provide a final recommended action by itself
• Generalize to unrelated image domains
• Perform object detection or image segmentation

Quick Start

Installation

Install the required libraries:

pip install torch transformers huggingface_hub pillow

Downloading the Model

The model files can be downloaded directly from Hugging Face Hub using hf_hub_download.

from huggingface_hub import hf_hub_download

model_path = hf_hub_download(
    repo_id="hur03/capturemate-image-classifier-v1-5class",
    filename="image_classifier.pt"
)

print(model_path)

You can also download the entire repository:

from huggingface_hub import snapshot_download

model_dir = snapshot_download(
    repo_id="hur03/capturemate-image-classifier-v1-5class"
)

print(model_dir)

Loading the Model

This model uses a custom PyTorch classifier class from the CaptureMate AI implementation.

It is not a standard AutoModelForImageClassification checkpoint.

The model architecture must match the architecture used during training before loading image_classifier.pt.

The general loading process is:

import torch
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(
    repo_id="hur03/capturemate-image-classifier-v1-5class",
    filename="image_classifier.pt"
)

# Import the custom classifier implementation
# from your CaptureMate AI codebase.
from your_module import CLIPUIClassifier

model = CLIPUIClassifier(num_labels=5)

state_dict = torch.load(
    model_path,
    map_location="cpu"
)

model.load_state_dict(state_dict)
model.eval()

Replace:

from your_module import CLIPUIClassifier

with the actual import path of the CLIPUIClassifier implementation from the CaptureMate AI codebase.

Image Preprocessing

The model is based on:

openai/clip-vit-base-patch32

Images should therefore be processed using the corresponding CLIP image processor.

from transformers import CLIPImageProcessor

processor = CLIPImageProcessor.from_pretrained(
    "openai/clip-vit-base-patch32"
)

Example:

from PIL import Image
from transformers import CLIPImageProcessor

processor = CLIPImageProcessor.from_pretrained(
    "openai/clip-vit-base-patch32"
)

image = Image.open("example_screenshot.jpg").convert("RGB")

inputs = processor(
    images=image,
    return_tensors="pt"
)

pixel_values = inputs["pixel_values"]

Inference

Once the custom CLIPUIClassifier architecture is available and the checkpoint has been loaded, inference can be performed with PyTorch.

import torch
from PIL import Image
from transformers import CLIPImageProcessor
from huggingface_hub import hf_hub_download

from your_module import CLIPUIClassifier

MODEL_ID = "hur03/capturemate-image-classifier-v1-5class"

LABELS = [
    "schedule",
    "shopping",
    "place",
    "memo",
    "unknown"
]

# Download checkpoint
model_path = hf_hub_download(
    repo_id=MODEL_ID,
    filename="image_classifier.pt"
)

# Load image processor
processor = CLIPImageProcessor.from_pretrained(
    "openai/clip-vit-base-patch32"
)

# Initialize model
model = CLIPUIClassifier(num_labels=5)

state_dict = torch.load(
    model_path,
    map_location="cpu"
)

model.load_state_dict(state_dict)
model.eval()

# Load screenshot
image = Image.open("example_screenshot.jpg").convert("RGB")

inputs = processor(
    images=image,
    return_tensors="pt"
)

# Inference
with torch.no_grad():
    outputs = model(inputs["pixel_values"])

    # Adjust this line if the custom model returns
    # an object containing logits instead of raw logits.
    logits = outputs

    probs = torch.softmax(logits, dim=-1)

pred_id = probs.argmax(dim=-1).item()
confidence = probs[0][pred_id].item()

result = {
    "label": LABELS[pred_id],
    "confidence": confidence
}

print(result)

Example output format:

{
    "label": "shopping",
    "confidence": 0.XX
}

The exact inference code may depend on the implementation of CLIPUIClassifier. The custom model architecture must match the architecture used when the checkpoint was trained.

Training Details

Training Data

The model was trained on screenshot images collected for the CaptureMate project.

The same screenshot category structure used by the OCR text classifier was used for image classification.

Dataset split:

The five target categories are:

Evaluation

Image-Only Performance

The image classifier was first evaluated independently on the CaptureMate test set.

The image-only classifier performs below the OCR text classifier, which supports its intended role as a complementary visual signal rather than the primary classifier.

Multimodal Performance

The image classifier was also evaluated as part of the CaptureMate multimodal classification pipeline.

The current fixed fusion uses:

Text weight  = 0.7
Image weight = 0.3

The resulting multimodal performance is:

The multimodal result combines predictions from both the OCR text classifier and this image classifier.

Therefore, the multimodal metrics above should not be interpreted as the standalone performance of this image model.

Limitations

The model is trained on a relatively small CaptureMate-specific screenshot dataset and has several important limitations.

Known limitations:

• It may not generalize well to screenshots outside the CaptureMate dataset distribution.
• Visually similar screenshots can be difficult to distinguish.
• Ambiguous unknown, shopping, and memo cases may be confused.
• Image-only performance is lower than the OCR text classifier.
• The model always produces one of the five class predictions.
• Visual appearance may vary significantly depending on application, operating system, theme, language, and screenshot layout.
• Screenshots containing primarily textual information may provide stronger signals to the OCR text classifier than to this image model.
• Confidence scores should not be interpreted as guaranteed probabilities of correctness.

Recommendations

This model is recommended as a complementary signal within a multimodal screenshot classification system.

For CaptureMate, the preferred usage is:

OCR Text
    ↓
Text Classifier
    ↓
Text Scores ─────────────┐
                         │
                         v
                    Score Fusion
                         ^
                         │
Image Scores ────────────┘
    ↑
Image Classifier
    ↑
Screenshot

The current CaptureMate implementation gives more weight to the OCR text classifier:

Text  = 70%
Image = 30%

This reflects the stronger standalone performance of the OCR text model on the current dataset.

The fusion weights may be adjusted in future versions based on validation experiments and additional training data.

Files

This repository contains the files required to distribute the trained image classifier.

image_classifier.pt

PyTorch checkpoint containing the trained CLIP vision encoder and custom classifier head weights.

preprocessor_config.json

Configuration for image preprocessing based on the CLIP image processor.

Because the model uses a custom PyTorch architecture, the checkpoint alone does not define the complete Python model class.

The corresponding CLIPUIClassifier implementation is required to reconstruct the architecture before loading the weights.

Technical Specifications

Model Architecture

Screenshot Image
       |
       v
CLIP Image Processor
       |
       v
CLIP Vision Encoder
(openai/clip-vit-base-patch32)
       |
       v
Visual Representation
       |
       v
Custom Classification Head
       |
       v
5-Class Logits
       |
       v
Softmax
       |
       v
Category Prediction

Model Objective

• Task: Single-label multi-class classification
• Input: Screenshot image
• Output: One of five CaptureMate categories
• Number of classes: 5
• Base encoder: CLIP ViT-B/32
• Framework: PyTorch

Label Mapping

0 -> schedule
1 -> shopping
2 -> place
3 -> memo
4 -> unknown

Relationship to the CaptureMate Text Classifier

CaptureMate uses two complementary classification models.

OCR Text Classifier

hur03/capturemate-category-classifier-v1-5class

The text classifier analyzes OCR-extracted screenshot text using a fine-tuned KLUE-RoBERTa model.

Image Classifier

hur03/capturemate-image-classifier-v1-5class

This repository contains the image classifier, which analyzes screenshot visual information using a CLIP-based architecture.

Together, the two models provide complementary information:

KLUE-RoBERTa
OCR/Text understanding
        +
CLIP
Visual understanding
        ↓
Multimodal screenshot classification

The current CaptureMate fusion strategy combines the two model outputs using fixed weighted score fusion.

Software

The model was developed using:

• PyTorch
• transformers
• huggingface_hub
• Pillow
• scikit-learn

Model Card Authors

CaptureMate