Upload README.md with huggingface_hub

README.md

@@ -2,107 +2,59 @@
 license: apache-2.0
 tags:
 - image-classification
-- checkbox-detection
-- computer-vision
+- transformers
 - pytorch
 datasets:
 - wendys-llc/chkbx
-metrics:
-- accuracy
-library_name: pytorch
 ---
 
-# Checkbox State Classifier
+# Checkbox Classifier
 
-This model classifies whether a checkbox is checked or 
-unchecked.
+Binary classifier for checkbox states (checked/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
+## Usage with Transformers
 
 ```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')
+from transformers import pipeline
 
-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])
-])
+# Load pipeline
+classifier = pipeline("image-classification", 
+                     model="wendys-llc/checkbox-classifier", 
+                     trust_remote_code=True)
 
 # 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
+from PIL import Image
+image = Image.open("checkbox.jpg")
+result = classifier(image)
+print(result)
+# [
+#   {'label': 'checked', 'score': 0.99},
+#   {'label': 'unchecked', 'score': 0.01}
+# ]
+```
 
-# Example usage
-result, conf = predict("checkbox.jpg")
-print(f"Result: {result} (confidence: {conf:.1%})")
+## Direct Usage
 
-Training Dataset
+```python
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+import torch
+from PIL import Image
+
+model = AutoModelForImageClassification.from_pretrained(
+    "wendys-llc/checkbox-classifier", 
+    trust_remote_code=True
+)
+processor = AutoImageProcessor.from_pretrained("wendys-llc/checkbox-classifier")
 
-Trained on https://huggingface.co/datasets/wendys-llc/chkbx
-dataset containing ~6,000 annotated checkbox images.
+image = Image.open("checkbox.jpg")
+inputs = processor(images=image, return_tensors="pt")
 
-Limitations
+with torch.no_grad():
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class = logits.argmax(-1).item()
+    
+print(model.config.id2label[predicted_class])
+```
 
-- 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
+## Accuracy: 97.1%