File size: 3,019 Bytes
af65138 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
---
license: apache-2.0
tags:
- image-classification
- checkbox-detection
- computer-vision
- pytorch
datasets:
- wendys-llc/chkbx
metrics:
- accuracy
library_name: pytorch
---
# Checkbox State Classifier
This model classifies whether a checkbox is checked or
unchecked.
## Model Details
- **Architecture**: EfficientNetV2-S (PyTorch)
- **Input Size**: 128x128 RGB images
- **Output**: Binary classification (unchecked: 0, checked: 1)
- **Validation Accuracy**: 97.1%
- **Training**: Mixed precision on A100 GPU
## Quick Start
```python
import torch
from PIL import Image
from torchvision import transforms
from huggingface_hub import hf_hub_download
import torch.nn as nn
from torchvision.models import efficientnet_v2_s,
EfficientNet_V2_S_Weights
# Define model architecture
class EfficientNetV2Classifier(nn.Module):
def __init__(self, num_classes=2, dropout_rate=0.3):
super().__init__()
self.backbone = efficientnet_v2_s(weights=EfficientNet
_V2_S_Weights.IMAGENET1K_V1)
num_features = self.backbone.classifier[1].in_features
self.backbone.classifier = nn.Sequential(
nn.Dropout(dropout_rate),
nn.Linear(num_features, 512),
nn.SiLU(inplace=True),
nn.BatchNorm1d(512),
nn.Dropout(dropout_rate),
nn.Linear(512, 256),
nn.SiLU(inplace=True),
nn.BatchNorm1d(256),
nn.Dropout(dropout_rate/2),
nn.Linear(256, num_classes)
)
def forward(self, x):
return self.backbone(x)
# Download and load model
model_path = hf_hub_download(repo_id="wendys-llc/checkbox-classifier",
filename="checkbox_classifier.pth")
checkpoint = torch.load(model_path, map_location='cpu')
model = EfficientNetV2Classifier(num_classes=2)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
# Image preprocessing
transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# Predict
def predict(image_path):
image = Image.open(image_path).convert('RGB')
input_tensor = transform(image).unsqueeze(0)
with torch.no_grad():
output = model(input_tensor)
probabilities = torch.nn.functional.softmax(output,
dim=1)
predicted = torch.argmax(probabilities, dim=1).item()
confidence = probabilities[0][predicted].item()
labels = {0: "unchecked", 1: "checked"}
return labels[predicted], confidence
# Example usage
result, conf = predict("checkbox.jpg")
print(f"Result: {result} (confidence: {conf:.1%})")
Training Dataset
Trained on https://huggingface.co/datasets/wendys-llc/chkbx
dataset containing ~6,000 annotated checkbox images.
Limitations
- Trained specifically on UI checkboxes, may not work well on
hand-drawn checkmarks
- Best performance on clear, high-contrast checkbox images
- Input images are resized to 128x128, very small checkboxes
may lose detail
|