Upload folder using huggingface_hub

README.md

@@ -1,270 +1,277 @@
-# Cervical Type Classification - Model Training
-
-## Overview
-
-This project classifies cervical images into 3 transformation zone types:
-- **Type 1**: Fully visible squamocolumnar junction (SCJ)
-- **Type 2**: Partially visible SCJ
-- **Type 3**: SCJ not visible (inside cervical canal)
-
-## Best Model Summary
-
-| Metric | Value |
-|--------|-------|
-| **Validation Accuracy** | **65.52%** |
-| **Macro F1** | **65.61%** |
-| Best Epoch | 34 |
-| Total Parameters | 1,327,235 |
-
+---
+license: mit
+tags:
+- image-classification
+- medical-imaging
+- cervical-cancer
+- pytorch
+- safetensors
+- cnn
+datasets:
+- custom
+metrics:
+- accuracy
+- f1
+pipeline_tag: image-classification
+library_name: pytorch
 ---
 
-## Best Model Configuration
-
-**Run Name:** `L32_64_128_256_Res_SE_lr5e-04_d0.3`
-
-### Architecture
+# CerviGuard - Cervical Transformation Zone Classifier
 
-| Component | Value |
-|-----------|-------|
-| Conv Layers | [32, 64, 128, 256] |
-| FC Layers | [256, 128] |
-| Kernel Size | 3x3 |
-| Pooling | MaxPool 2x2 |
-| Batch Normalization | Yes |
-| Activation | ReLU |
-| Residual Connections | **Yes** |
-| SE Attention | **Yes** |
+## Model Description
 
-### Training Settings
+This model classifies cervical images into 3 transformation zone types, which is important for colposcopy evaluation and cervical cancer screening.
 
-| Parameter | Value |
-|-----------|-------|
-| Learning Rate | 5e-4 |
-| Weight Decay | 1e-4 |
-| Dropout | 0.3 |
-| Batch Size | 32 |
-| Focal Loss Gamma | 2.0 |
-| Label Smoothing | 0.1 |
-| Data Augmentation | Yes |
+| Label | Type | Description |
+|-------|------|-------------|
+| 0 | Type 1 | Transformation zone fully visible on ectocervix |
+| 1 | Type 2 | Transformation zone partially visible (extends into endocervical canal) |
+| 2 | Type 3 | Transformation zone not visible (entirely within endocervical canal) |
 
 ---
 
-## Performance Metrics
-
-### Per-Class Metrics
+## Model Architecture
 
-| Class | Precision | Recall | F1-Score | Support |
-|-------|-----------|--------|----------|---------|
-| **Type 1** | 79.26% | 61.49% | 69.26% | 348 |
-| **Type 2** | 58.09% | 75.29% | 65.58% | 348 |
-| **Type 3** | 64.40% | 59.77% | 62.00% | 348 |
-| **Macro Avg** | 67.25% | 65.52% | **65.61%** | 1044 |
+### Overview
 
-### Confusion Matrix
+**BaseCNN** - A simple convolutional neural network with 4 conv blocks and 2 fully connected layers.
 
 ```
-                 Predicted
-              Type 1  Type 2  Type 3
-Actual Type 1   214      84      50
-       Type 2    21     262      65
-       Type 3    35     105     208
+┌─────────────────────────────────────────────────────────────┐
+│                    INPUT (256×256×3)                        │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│  CONV BLOCK 1                                               │
+│  Conv2d(3→32, 3×3) → BatchNorm2d → ReLU → MaxPool2d(2×2)   │
+│  Output: 128×128×32                                         │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│  CONV BLOCK 2                                               │
+│  Conv2d(32→64, 3×3) → BatchNorm2d → ReLU → MaxPool2d(2×2)  │
+│  Output: 64×64×64                                           │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│  CONV BLOCK 3                                               │
+│  Conv2d(64→128, 3×3) → BatchNorm2d → ReLU → MaxPool2d(2×2) │
+│  Output: 32×32×128                                          │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│  CONV BLOCK 4                                               │
+│  Conv2d(128→256, 3×3) → BatchNorm2d → ReLU → MaxPool2d(2×2)│
+│  Output: 16×16×256                                          │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│  GLOBAL POOLING                                             │
+│  AdaptiveAvgPool2d(1×1)                                     │
+│  Output: 1×1×256 → Flatten → 256                            │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│  FC BLOCK 1                                                 │
+│  Linear(256→256) → ReLU → Dropout(0.4)                     │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│  FC BLOCK 2                                                 │
+│  Linear(256→128) → ReLU → Dropout(0.4)                     │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│  CLASSIFIER                                                 │
+│  Linear(128→3)                                              │
+└─────────────────────────────────────────────────────────────┘
+                            │
+┌─────────────────────────────────────────────────────────────┐
+│                 OUTPUT (3 logits)                           │
+│            [Type 1, Type 2, Type 3]                         │
+└─────────────────────────────────────────────────────────────┘
 ```
 
-### Interpretation
-
-| Finding | Implication |
-|---------|-------------|
-| Type 1 has highest precision (79%) | When model predicts Type 1, it's usually correct |
-| Type 2 has highest recall (75%) | Model catches most Type 2 cases |
-| Type 3 has lowest metrics | Hardest to classify - often confused with Type 2 |
-| Type 2 ↔ Type 3 confusion is common | 105 Type 3 misclassified as Type 2 |
+### Layer Details
+
+| Layer | Type | In Channels | Out Channels | Kernel | Output Size |
+|-------|------|-------------|--------------|--------|-------------|
+| conv_layers.0 | Conv2d | 3 | 32 | 3×3 | 256×256×32 |
+| conv_layers.1 | BatchNorm2d | 32 | 32 | - | 256×256×32 |
+| conv_layers.2 | ReLU | - | - | - | 256×256×32 |
+| conv_layers.3 | MaxPool2d | - | - | 2×2 | 128×128×32 |
+| conv_layers.4 | Conv2d | 32 | 64 | 3×3 | 128×128×64 |
+| conv_layers.5 | BatchNorm2d | 64 | 64 | - | 128×128×64 |
+| conv_layers.6 | ReLU | - | - | - | 128×128×64 |
+| conv_layers.7 | MaxPool2d | - | - | 2×2 | 64×64×64 |
+| conv_layers.8 | Conv2d | 64 | 128 | 3×3 | 64×64×128 |
+| conv_layers.9 | BatchNorm2d | 128 | 128 | - | 64×64×128 |
+| conv_layers.10 | ReLU | - | - | - | 64×64×128 |
+| conv_layers.11 | MaxPool2d | - | - | 2×2 | 32×32×128 |
+| conv_layers.12 | Conv2d | 128 | 256 | 3×3 | 32×32×256 |
+| conv_layers.13 | BatchNorm2d | 256 | 256 | - | 32×32×256 |
+| conv_layers.14 | ReLU | - | - | - | 32×32×256 |
+| conv_layers.15 | MaxPool2d | - | - | 2×2 | 16×16×256 |
+| adaptive_pool | AdaptiveAvgPool2d | - | - | - | 1×1×256 |
+| fc_layers.0 | Linear | 256 | 256 | - | 256 |
+| fc_layers.1 | ReLU | - | - | - | 256 |
+| fc_layers.2 | Dropout | - | - | p=0.4 | 256 |
+| fc_layers.3 | Linear | 256 | 128 | - | 128 |
+| fc_layers.4 | ReLU | - | - | - | 128 |
+| fc_layers.5 | Dropout | - | - | p=0.4 | 128 |
+| classifier | Linear | 128 | 3 | - | 3 |
+
+### Model Summary
+
+| Property | Value |
+|----------|-------|
+| **Total Parameters** | 488,451 |
+| **Trainable Parameters** | 488,451 |
+| **Input Size** | (B, 3, 256, 256) |
+| **Output Size** | (B, 3) |
+| **Model Size** | ~1.9 MB |
 
 ---
 
-## Grid Search Results
-
-A grid search of 32 configurations was performed on January 17, 2026.
-
-### Search Space
-
-| Parameter | Values Tested |
-|-----------|---------------|
-| Conv Layers | [32,64,128,256], [64,128,256] |
-| Learning Rate | 5e-4, 1e-4 |
-| Dropout | 0.3, 0.4 |
-| Residual | Yes, No |
-| SE Attention | Yes, No |
-
-### Top 10 Configurations
-
-| Rank | Configuration | Accuracy | Key Features |
-|------|--------------|----------|--------------|
-| 1 | L32_64_128_256_Res_SE_lr5e-04_d0.3 | **65.52%** | 4-layer, Res+SE |
-| 2 | L64_128_256_Res_SE_lr5e-04_d0.3 | 65.04% | 3-layer, Res+SE |
-| 3 | L32_64_128_256_Res_SE_lr1e-04_d0.3 | 64.94% | 4-layer, lower LR |
-| 4 | L64_128_256_Res_SE_lr1e-04_d0.3 | 64.37% | 3-layer, lower LR |
-| 5 | L32_64_128_256_Res_SE_lr5e-04_d0.4 | 64.18% | Higher dropout |
-| 6 | L32_64_128_256_Res_lr5e-04_d0.4 | 64.08% | No SE |
-| 7 | L32_64_128_256_Res_lr1e-04_d0.3 | 63.60% | No SE, lower LR |
-| 8 | L32_64_128_256_Res_SE_lr1e-04_d0.4 | 63.51% | Lower LR, higher dropout |
-| 9 | L64_128_256_Res_SE_lr5e-04_d0.4 | 63.22% | 3-layer, higher dropout |
-| 10 | L64_128_256_Res_SE_lr1e-04_d0.4 | 63.12% | 3-layer, lower LR |
-
-### Key Findings
-
-| Finding | Evidence |
-|---------|----------|
-| **Residual + SE is critical** | Top 10 models all use residual connections; top 4 use both Res+SE |
-| **4-layer network is better** | [32,64,128,256] outperforms [64,128,256] |
-| **Higher LR (5e-4) preferred** | 5e-4 consistently beats 1e-4 |
-| **Lower dropout (0.3) preferred** | 0.3 dropout outperforms 0.4 |
-| **Plain CNN performs worst** | Models without Res or SE are at the bottom |
-
-### What Worked vs What Didn't
-
-| Worked | Didn't Work |
-|--------|-------------|
-| Residual connections | Plain convolutions |
-| SE attention blocks | No attention |
-| 4 conv layers | 3 conv layers |
-| LR = 5e-4 | LR = 1e-4 (too slow) |
-| Dropout = 0.3 | Dropout = 0.4 (too aggressive) |
-| Focal Loss | - |
-| Label smoothing 0.1 | - |
+## Training Configuration
+
+| Parameter | Value |
+|-----------|-------|
+| Learning Rate | 1e-4 |
+| Batch Size | 32 |
+| Dropout | 0.4 |
+| Optimizer | Adam |
+| Loss Function | CrossEntropyLoss |
+| Epochs | 50 |
+| Best Epoch | 41 |
 
 ---
 
-## Data
+## Performance
 
-| Split | Samples | Classes | Distribution |
-|-------|---------|---------|--------------|
-| Train | ~7,000 | 3 | Balanced after augmentation |
-| Test | 1,044 | 3 | [348, 348, 348] |
+| Metric | Value |
+|--------|-------|
+| **Validation Accuracy** | 61.69% |
+| **Macro F1 Score** | 61.81% |
 
-### Image Specifications
+### Per-Class Performance
 
-- Size: Variable (resized during training)
-- Channels: 3 (RGB)
-- Source: Colposcopy images
+| Type | Precision | Recall | F1 Score |
+|------|-----------|--------|----------|
+| Type 1 | - | - | 68.32% |
+| Type 2 | - | - | 56.41% |
+| Type 3 | - | - | 60.69% |
 
 ---
 
-## Model Files
-
-### Best Model Location
+## Usage
 
-```
-./best_model.pth  (this folder)
-```
+### Installation
 
-Original training output:
-```
-/data/downloads/cervical_type/_output/grid_search_v2_20260117_212011/run_001_L32_64_128_256_Res_SE_lr5e-04_d0.3/
+```bash
+pip install torch torchvision safetensors huggingface_hub
 ```
 
-### Checkpoint Contents
+### Quick Start (Local)
 
 ```python
-{
-    "epoch": 34,
-    "model_state_dict": ...,
-    "optimizer_state_dict": ...,
-    "scheduler_state_dict": ...,
-    "metrics": {...},
-    "model_config": {...}
-}
-```
+import torch
+from PIL import Image
+from torchvision import transforms
 
-### Files in This Folder
+# Load model
+from model import BaseCNN
+model = BaseCNN.from_pretrained("./")
+model.eval()
 
-| File | Description |
-|------|-------------|
-| `best_model.pth` | Model checkpoint (weights + optimizer state) |
-| `config.json` | Training configuration used |
-| `training_history.json` | Loss/accuracy per epoch |
-| `grid_search_summary.json` | All 32 grid search results |
-| `README.md` | This file |
+# Preprocess image
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+])
 
-### Loading the Model
+image = Image.open("cervical_image.jpg").convert("RGB")
+input_tensor = transform(image).unsqueeze(0)
 
-```python
-import torch
+# Inference
+with torch.no_grad():
+    output = model(input_tensor)
+    probabilities = torch.softmax(output, dim=1)
+    prediction = output.argmax(dim=1).item()
 
-# Load checkpoint (from this folder)
-checkpoint = torch.load('best_model.pth', weights_only=False)
-
-# Create model with same config
-model = BaseCNN(
-    conv_layers=[32, 64, 128, 256],
-    fc_layers=[256, 128],
-    num_classes=3,
-    dropout=0.3,
-    use_residual=True,
-    use_se_attention=True
-)
-
-# Load weights
-model.load_state_dict(checkpoint['model_state_dict'])
-model.eval()
+labels = ["Type 1", "Type 2", "Type 3"]
+print(f"Prediction: {labels[prediction]}")
+print(f"Confidence: {probabilities[0][prediction]:.2%}")
 ```
 
----
+### Load from Hugging Face Hub
 
-## Output Structure
+```python
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+import torch
+import json
+import importlib.util
+
+# Download files
+repo_id = "toderian/cerviguard_transfer_zones"
+model_weights = hf_hub_download(repo_id, "model.safetensors")
+config_file = hf_hub_download(repo_id, "config.json")
+model_file = hf_hub_download(repo_id, "model.py")
+
+# Load model class dynamically
+spec = importlib.util.spec_from_file_location("model", model_file)
+model_module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(model_module)
+
+# Load config and create model
+with open(config_file) as f:
+    config = json.load(f)
+
+model = model_module.BaseCNN(**config['model_config'])
+model.load_state_dict(load_file(model_weights))
+model.eval()
 
-```
-_output/
-└── grid_search_v2_20260117_212011/
-    ├── grid_search_config.json    # Search space definition
-    ├── all_results.json           # All 32 run results
-    ├── summary.json               # Sorted results + best run
-    ├── logs/
-    │   └── grid_search.log
-    └── run_001_.../               # Best run
-        ├── checkpoints/
-        │   ├── best_model.pth     # Best validation accuracy
-        │   ├── latest.pth         # Final epoch
-        │   └── epoch_*.pth        # Periodic saves
-        └── logs/
-            ├── run_config.json
-            └── training_history.json
+# Now use model for inference
 ```
 
 ---
 
-## Comparison with v1 Baseline
-
-| Version | Accuracy | Improvement |
-|---------|----------|-------------|
-| v1 Baseline | 61.69% | - |
-| **v2 Best (Res+SE)** | **65.52%** | **+3.83%** |
+## Files in This Repository
 
-The addition of residual connections and SE attention improved accuracy by nearly 4%.
+| File | Description |
+|------|-------------|
+| `model.safetensors` | Model weights (SafeTensors format, recommended) |
+| `pytorch_model.bin` | Model weights (PyTorch format, backup) |
+| `config.json` | Model architecture configuration |
+| `model.py` | Model class definition (BaseCNN) |
+| `preprocessor_config.json` | Image preprocessing configuration |
+| `README.md` | This model card |
 
 ---
 
-## Recommendations for Future Work
+## Limitations
 
-1. **Try deeper networks** - Add 5th conv layer [32, 64, 128, 256, 512]
-2. **Transfer learning** - Use pretrained EfficientNet or ResNet backbone
-3. **Address Type 3 confusion** - Type 3 is often misclassified as Type 2
-4. **Ensemble methods** - Combine top 3-5 models
-5. **Test Time Augmentation** - Average predictions over augmented versions
-6. **More training data** - Current ~7k samples may be limiting
+- Model was trained on a specific dataset and may not generalize to all cervical imaging equipment
+- Type 2 classification has lower accuracy (56.41% F1) as it represents an intermediate state
+- Input images should be 256×256 RGB
+- This is a custom PyTorch model, not compatible with `transformers.AutoModel`
 
 ---
 
-## Quick Start
-
-```bash
-# Run the best configuration
-python train_grid_v2.py
+## Citation
 
-# Or load and evaluate the best model
-python evaluate.py --model /path/to/best_model.pth
+```bibtex
+@misc{cerviguard-transfer-zones,
+  title={CerviGuard Cervical Transformation Zone Classifier},
+  author={toderian},
+  year={2026},
+  howpublished={\url{https://huggingface.co/toderian/cerviguard_transfer_zones}}
+}
 ```
 
 ---
 
-*Last updated: January 2026*
-*Grid search: 32 configurations, ~15 hours on single GPU*
+## License
+
+MIT License

config.json

@@ -1,26 +1,36 @@
 {
-  "batch_size": 32,
-  "learning_rate": 0.0005,
-  "weight_decay": 0.0001,
-  "layers": [
-    32,
-    64,
-    128,
-    256
-  ],
-  "use_residual": true,
-  "use_se_attention": true,
-  "focal_gamma": 2.0,
-  "label_smoothing": 0.1,
-  "dropout": 0.3,
-  "kernel": 3,
-  "batchnorm": true,
-  "activation": "ReLU",
-  "pool": true,
-  "fc_multipliers": [
-    1.0,
-    0.5
-  ],
-  "nr_classes": 3,
-  "augmentation": true
+  "model_type": "BaseCNN",
+  "model_config": {
+    "layers": [
+      32,
+      64,
+      128,
+      256
+    ],
+    "kernel": 3,
+    "padding": 1,
+    "stride": 1,
+    "batchnorm": true,
+    "bn_pre_activ": true,
+    "activation": "ReLU",
+    "dropout": 0.4,
+    "pool": true,
+    "fc_layers": [
+      256,
+      128
+    ],
+    "nr_classes": 3,
+    "in_channels": 3
+  },
+  "num_labels": 3,
+  "id2label": {
+    "0": "Type 1",
+    "1": "Type 2",
+    "2": "Type 3"
+  },
+  "label2id": {
+    "Type 1": 0,
+    "Type 2": 1,
+    "Type 3": 2
+  }
 }

model.py

@@ -0,0 +1,244 @@
+"""
+Cervical Type Classification Model
+
+This module contains the BaseCNN model for classifying cervical images
+into 3 transformation zone types.
+
+Usage:
+    from model import BaseCNN
+
+    # Load pretrained model
+    model = BaseCNN.from_pretrained("./")
+
+    # Or create from scratch
+    model = BaseCNN(
+        layers=[32, 64, 128, 256],
+        fc_layers=[256, 128],
+        nr_classes=3
+    )
+"""
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+try:
+    from safetensors.torch import load_file, save_file
+    HAS_SAFETENSORS = True
+except ImportError:
+    HAS_SAFETENSORS = False
+
+
+class BaseCNN(nn.Module):
+    """
+    Simple CNN for cervical type classification.
+
+    Classifies cervical images into 3 transformation zone types:
+    - Type 1: Transformation zone fully visible on ectocervix
+    - Type 2: Transformation zone partially visible
+    - Type 3: Transformation zone not visible (within endocervical canal)
+
+    Args:
+        layers: List of output channels for each conv layer. Default: [32, 64, 128, 256]
+        kernel: Kernel size for conv layers. Default: 3
+        padding: Padding for conv layers. Default: 1
+        stride: Stride for conv layers. Default: 1
+        batchnorm: Whether to use batch normalization. Default: True
+        bn_pre_activ: Whether to apply BN before activation. Default: True
+        activation: Activation function name. Default: 'ReLU'
+        dropout: Dropout rate for FC layers. Default: 0.4
+        pool: Whether to use max pooling after each conv. Default: True
+        fc_layers: List of FC layer sizes. Default: [256, 128]
+        nr_classes: Number of output classes. Default: 3
+        in_channels: Number of input channels. Default: 3
+    """
+
+    def __init__(
+        self,
+        layers: list = None,
+        kernel: int = 3,
+        padding: int = 1,
+        stride: int = 1,
+        batchnorm: bool = True,
+        bn_pre_activ: bool = True,
+        activation: str = 'ReLU',
+        dropout: float = 0.4,
+        pool: bool = True,
+        fc_layers: list = None,
+        nr_classes: int = 3,
+        in_channels: int = 3,
+    ):
+        super().__init__()
+
+        # Store config for serialization
+        self.config = {
+            'layers': layers or [32, 64, 128, 256],
+            'kernel': kernel,
+            'padding': padding,
+            'stride': stride,
+            'batchnorm': batchnorm,
+            'bn_pre_activ': bn_pre_activ,
+            'activation': activation,
+            'dropout': dropout,
+            'pool': pool,
+            'fc_layers': fc_layers or [256, 128],
+            'nr_classes': nr_classes,
+            'in_channels': in_channels,
+        }
+
+        layers = self.config['layers']
+        fc_layers = self.config['fc_layers']
+
+        # Activation function
+        activation_fn = getattr(nn, activation)
+
+        # Build convolutional layers (ModuleList to match original)
+        self.conv_layers = nn.ModuleList()
+        prev_channels = in_channels
+
+        for out_channels in layers:
+            self.conv_layers.append(
+                nn.Conv2d(prev_channels, out_channels, kernel, stride, padding)
+            )
+            if batchnorm and bn_pre_activ:
+                self.conv_layers.append(nn.BatchNorm2d(out_channels))
+            self.conv_layers.append(activation_fn())
+            if batchnorm and not bn_pre_activ:
+                self.conv_layers.append(nn.BatchNorm2d(out_channels))
+            if pool:
+                self.conv_layers.append(nn.MaxPool2d(2, 2))
+            prev_channels = out_channels
+
+        # Global average pooling
+        self.adaptive_pool = nn.AdaptiveAvgPool2d(1)
+
+        # Build fully connected layers (ModuleList to match original)
+        self.fc_layers = nn.ModuleList()
+        prev_features = layers[-1]
+
+        for fc_size in fc_layers:
+            self.fc_layers.append(nn.Linear(prev_features, fc_size))
+            self.fc_layers.append(activation_fn())
+            self.fc_layers.append(nn.Dropout(dropout))
+            prev_features = fc_size
+
+        # Final classifier (separate, to match original)
+        self.classifier = nn.Linear(prev_features, nr_classes)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x: Input tensor of shape (batch_size, 3, 256, 256)
+
+        Returns:
+            Logits tensor of shape (batch_size, num_classes)
+        """
+        for layer in self.conv_layers:
+            x = layer(x)
+
+        x = self.adaptive_pool(x)
+        x = x.view(x.size(0), -1)
+
+        for layer in self.fc_layers:
+            x = layer(x)
+
+        x = self.classifier(x)
+        return x
+
+    @classmethod
+    def from_pretrained(cls, model_path: str, device: str = 'cpu') -> 'BaseCNN':
+        """
+        Load a pretrained model from a directory.
+
+        Args:
+            model_path: Path to directory containing model files
+            device: Device to load model on ('cpu' or 'cuda')
+
+        Returns:
+            Loaded model in eval mode
+        """
+        model_path = Path(model_path)
+
+        # Load config
+        config_path = model_path / 'config.json'
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+
+        # Create model
+        model = cls(**config['model_config'])
+
+        # Load weights (prefer safetensors)
+        safetensors_path = model_path / 'model.safetensors'
+        pytorch_path = model_path / 'pytorch_model.bin'
+
+        if safetensors_path.exists() and HAS_SAFETENSORS:
+            state_dict = load_file(str(safetensors_path), device=device)
+        elif pytorch_path.exists():
+            state_dict = torch.load(pytorch_path, map_location=device, weights_only=True)
+        else:
+            raise FileNotFoundError(f"No model weights found in {model_path}")
+
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+        return model
+
+    def save_pretrained(self, save_path: str) -> None:
+        """
+        Save model in Hugging Face compatible format.
+
+        Args:
+            save_path: Directory to save model files
+        """
+        save_path = Path(save_path)
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        # Save config
+        config = {
+            'model_type': 'BaseCNN',
+            'model_config': self.config,
+            'num_labels': self.config['nr_classes'],
+            'id2label': {
+                '0': 'Type 1',
+                '1': 'Type 2',
+                '2': 'Type 3'
+            },
+            'label2id': {
+                'Type 1': 0,
+                'Type 2': 1,
+                'Type 3': 2
+            }
+        }
+        with open(save_path / 'config.json', 'w') as f:
+            json.dump(config, f, indent=2)
+
+        # Save weights
+        state_dict = {k: v.contiguous() for k, v in self.state_dict().items()}
+
+        # SafeTensors format (recommended)
+        if HAS_SAFETENSORS:
+            save_file(state_dict, str(save_path / 'model.safetensors'))
+
+        # PyTorch format (backup)
+        torch.save(state_dict, save_path / 'pytorch_model.bin')
+
+
+# Label mappings
+ID2LABEL = {0: 'Type 1', 1: 'Type 2', 2: 'Type 3'}
+LABEL2ID = {'Type 1': 0, 'Type 2': 1, 'Type 3': 2}
+
+
+if __name__ == '__main__':
+    # Quick test
+    model = BaseCNN()
+    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Test forward pass
+    x = torch.randn(1, 3, 256, 256)
+    y = model(x)
+    print(f"Input shape: {x.shape}")
+    print(f"Output shape: {y.shape}")

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:beb3e17da6b94596232aa18078b9d22872f4711c7c1ef21a35f3277175d14063
+size 1960588

preprocessor_config.json

@@ -0,0 +1,14 @@
+{
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [0.5, 0.5, 0.5],
+  "image_std": [0.5, 0.5, 0.5],
+  "image_processor_type": "ImageProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 256,
+    "width": 256
+  }
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d88b2345cde9dcdc7fc2b8ba76edb2c64abfbc274f320bd55ad9e12801c9b00
+size 1969453