Update README.md

README.md

@@ -1,3 +1,284 @@
----
-license: mit
----
+---
+license: mit
+datasets:
+- ArtifactClfDurham/OrientalMuseum-white
+language:
+- en
+base_model:
+- google/efficientnet-b0
+tags:
+- artifact
+- museum
+---
+
+# Artifact Classification Model v2 - Best Model Usage Guide
+
+This directory contains the improved v2 artifact classification model with state-of-the-art performance for classifying museum artifacts by both object type and material.
+
+## Model Overview
+
+The v2 model is an advanced multi-output neural network that predicts two attributes simultaneously:
+- **Object Name**: The type/category of the artifact (e.g., "vase", "statue", "pottery")
+- **Material**: The material composition (e.g., "ceramic", "bronze", "stone")
+
+### Key Improvements Over v1
+- **EfficientNet Backbone**: Uses EfficientNet-B0 instead of ResNet-50 for better feature extraction
+- **Attention Mechanism**: Includes an attention layer to focus on relevant features
+- **Advanced Training**: Incorporates CutMix augmentation, Focal Loss, and mixed precision training
+- **Better Regularization**: Uses dropout and batch normalization for improved generalization
+
+## Quick Start
+
+### Prerequisites
+
+Ensure you have the required dependencies installed:
+
+```bash
+pip install torch>=2.0.0 torchvision>=0.15.0 datasets>=2.0.0 pillow>=9.0.0 timm>=1.0.22 huggingface-hub>=0.15.0
+```
+
+### Basic Inference
+
+```python
+import torch
+from PIL import Image
+from torchvision import transforms
+import sys
+import os
+
+# Add the project root to Python path
+sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
+
+from main import load_model, run_inference
+
+# Load the model
+model_path = "model/v2/best_model.pth"
+model, label_mappings = load_model(model_path)
+
+# Prepare image
+image_path = "path/to/your/artifact.jpg"
+image = Image.open(image_path).convert('RGB')
+
+# Preprocessing transform
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+pixel_values = transform(image).unsqueeze(0)  # Add batch dimension
+
+# Run inference
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+preds_obj, confs_obj, preds_mat, confs_mat = run_inference(model, pixel_values, device)
+
+# Get predictions
+object_pred_id = preds_obj[0].item()
+material_pred_id = preds_mat[0].item()
+object_conf = confs_obj[0].item()
+material_conf = confs_mat[0].item()
+
+# Convert IDs to labels
+object_name = label_mappings['object_name'].get(object_pred_id, f"class_{object_pred_id}")
+material_name = label_mappings['material'].get(material_pred_id, f"class_{material_pred_id}")
+
+print(f"Predicted Object: {object_name} (confidence: {object_conf:.3f})")
+print(f"Predicted Material: {material_name} (confidence: {material_conf:.3f})")
+```
+
+## Model Files
+
+- **`best_model.pth`**: The best performing model checkpoint with trained weights and label mappings
+- **`model_improved.pth`**: Final model after complete training
+- **`checkpoint_epoch_*.pth`**: Intermediate checkpoints saved during training
+- **`train.py`**: Training script used to create this model
+
+## Model Architecture
+
+```python
+ImprovedMultiOutputModel(
+    backbone: EfficientNet-B0 (pretrained)
+    attention: Linear(1280 → 512 → 1280) with Sigmoid
+    object_classifier: Linear(1280 → 1024 → 512 → num_object_classes)
+    material_classifier: Linear(1280 → 1024 → 512 → num_material_classes)
+)
+```
+
+### Input Requirements
+- **Image Size**: 224×224 pixels (automatically resized and cropped)
+- **Format**: RGB images
+- **Normalization**: ImageNet normalization (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+### Output Format
+Returns a dictionary with:
+- `'object_name'`: Logits for object classification
+- `'material'`: Logits for material classification
+
+## Evaluation
+
+### Using the Main Evaluation Script
+
+To evaluate the model on the Oriental Museum dataset:
+
+```bash
+# Evaluate on validation set
+python main.py --model_file model/v2/best_model.pth --output eval_results_v2.json
+
+# Evaluate with custom batch size
+python main.py --model_file model/v2/best_model.pth --batch_size 16 --output eval_results_v2.json
+```
+
+### Evaluation Metrics
+
+The evaluation script provides:
+- **Object Classification Accuracy**: Accuracy for object name prediction
+- **Material Classification Accuracy**: Accuracy for material prediction
+- **Overall Accuracy**: Samples where both predictions are correct
+- **Confidence Analysis**: Average confidence for correct vs incorrect predictions
+- **Per-sample Predictions**: Detailed results for each test sample
+
+### Expected Performance
+
+Based on validation during training:
+- Object Classification: ~85-90% accuracy
+- Material Classification: ~80-85% accuracy
+- Overall Accuracy: ~75-80% accuracy
+
+*Note: Actual performance may vary depending on the evaluation dataset and preprocessing.*
+
+## Training Details
+
+The model was trained with the following configuration:
+
+- **Dataset**: ArtifactClfDurham/OrientalMuseum-white
+- **Training Split**: 85% of data
+- **Validation Split**: 15% of data
+- **Batch Size**: 32
+- **Epochs**: 20
+- **Optimizer**: AdamW with differential learning rates
+  - Backbone: 2e-4 (0.1× base LR)
+  - Heads: 2e-3 (base LR)
+- **Augmentation**: Advanced (CutMix, rotation, color jitter, Gaussian blur)
+- **Loss Function**: Cross-Entropy (or Focal Loss if enabled)
+- **Scheduler**: Cosine annealing with warmup
+
+### Advanced Training Features
+
+- **CutMix Augmentation**: Randomly mixes image patches between samples
+- **Focal Loss**: Addresses class imbalance (optional)
+- **Mixed Precision**: Automatic mixed precision training for speed
+- **Gradient Scaling**: Prevents gradient underflow
+- **Early Stopping**: Saves best model based on validation accuracy
+
+## Usage Examples
+
+### Batch Inference
+
+```python
+import torch
+from PIL import Image
+from torchvision import transforms
+import sys
+import os
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
+from main import load_model, run_inference
+
+# Load model
+model, label_mappings = load_model("model/v2/best_model.pth")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load multiple images
+image_paths = ["artifact1.jpg", "artifact2.jpg", "artifact3.jpg"]
+images = []
+
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+for path in image_paths:
+    img = Image.open(path).convert('RGB')
+    images.append(transform(img))
+
+# Batch tensor
+batch = torch.stack(images)
+
+# Run inference
+preds_obj, confs_obj, preds_mat, confs_mat = run_inference(model, batch, device)
+
+# Process results
+for i, (obj_pred, obj_conf, mat_pred, mat_conf) in enumerate(zip(preds_obj, confs_obj, preds_mat, confs_mat)):
+    obj_name = label_mappings['object_name'].get(obj_pred.item(), f"class_{obj_pred.item()}")
+    mat_name = label_mappings['material'].get(mat_pred.item(), f"class_{mat_pred.item()}")
+
+    print(f"Image {i+1}:")
+    print(f"  Object: {obj_name} ({obj_conf:.3f})")
+    print(f"  Material: {mat_name} ({mat_conf:.3f})")
+```
+
+### Custom Dataset Evaluation
+
+```python
+from datasets import load_dataset
+from main import load_model
+import json
+
+# Load your custom dataset
+dataset = load_dataset("your-dataset", split="test")
+
+# Load model
+model, label_mappings = load_model("model/v2/best_model.pth")
+
+# Run evaluation (modify main.py evaluation logic as needed)
+# ... evaluation code ...
+```
+
+## Troubleshooting
+
+### Common Issues
+
+1. **CUDA Out of Memory**
+   - Reduce batch size: `--batch_size 8`
+   - Use CPU: Set device to "cpu"
+
+2. **Import Errors**
+   - Ensure all dependencies are installed
+   - Check Python path includes project root
+
+3. **Model Loading Errors**
+   - Verify the model file path is correct
+   - Ensure PyTorch version compatibility
+
+4. **Low Confidence Scores**
+   - Model may not be trained on similar artifacts
+   - Check image preprocessing matches training setup
+
+### Performance Tips
+
+- Use GPU for faster inference
+- Process images in batches for efficiency
+- Use the best_model.pth for production use
+- Consider model quantization for deployment
+
+## Model Limitations
+
+- Trained specifically on Oriental Museum artifacts
+- May not generalize well to artifacts from other cultures/regions
+- Performance depends on image quality and lighting
+- Multi-output nature may have trade-offs between object and material accuracy
+
+## Contributing
+
+To improve the model:
+1. Use the training script with different hyperparameters
+2. Experiment with different backbones
+3. Add more advanced augmentations
+4. Fine-tune on additional datasets
+
+## License
+
+This model is part of the artifact identification project. Check the main project license for usage terms.