File size: 2,868 Bytes
a05f853 6ac5a86 a05f853 6ac5a86 a05f853 6ac5a86 a05f853 6ac5a86 a05f853 6ac5a86 a05f853 6ac5a86 a05f853 6ac5a86 |
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 110 111 112 |
---
language: en
license: mit
library_name: transformers
pipeline_tag: image-classification
tags:
- vision
- damage-detection
- classification
- vit
- household-items
datasets:
- custom
---
# damage-classifier-multi-task
## Home Damage Classification Model
This model was trained to classify damage to household items, identifying both the item category, damage type, and damage severity.
### Model Type
- **Architecture**: Vision Transformer (ViT)
- **Classification Approach**: Multi_Task
- **Base Model**: google/vit-base-patch16-224
### Supported Categories
#### Item Categories
- Microwave
- Wall
- Window
- Fence
- Glass
- Fishbowl
#### Damage Types
- Scratch
- Dent
- Break
- Burn
- Water Damage
#### Severity Levels
- No Damage
- Minor Damage
- Moderate Damage
- Severe Damage
### Multi-Task Architecture
This model uses a multi-task learning approach with:
1. A shared Vision Transformer (ViT) backbone that extracts features from the input image
2. Separate classification heads for:
- Item category identification
- Damage type classification
- Damage severity assessment
This approach allows the model to share knowledge between related tasks while making separate predictions for each aspect.
#### Advantages of Multi-Task Learning
- Shares knowledge across related tasks
- Requires fewer examples per combination
- Can perform well even with missing combinations
- Independent predictions for each aspect
### Usage
```python
from transformers import ViTFeatureExtractor
from PIL import Image
import torch
# Load model and feature extractor
model = torch.load("pytorch_model.bin") # Or use your preferred loading method
feature_extractor = ViTFeatureExtractor.from_pretrained("USER/REPO_NAME")
# Prepare image
image = Image.open("path/to/image.jpg").convert("RGB")
inputs = feature_extractor(images=image, return_tensors="pt")
# Get predictions
outputs = model(**inputs)
# Process multi-task outputs
item_logits = outputs['item_logits']
damage_logits = outputs['damage_type_logits']
severity_logits = outputs['severity_logits']
# Get predicted classes
item_class = torch.argmax(item_logits, dim=1).item()
damage_class = torch.argmax(damage_logits, dim=1).item()
severity_class = torch.argmax(severity_logits, dim=1).item()
# Map to class names (replace with your class mappings)
item_categories = ["microwave", "wall", "window", "fence", "glass", "fishbowl"]
damage_types = ["scratch", "dent", "break", "burn", "water_damage"]
severity_levels = ["no_damage", "minor_damage", "moderate_damage", "severe_damage"]
print(f"Item: {item_categories[item_class]}")
print(f"Damage Type: {damage_types[damage_class]}")
print(f"Severity: {severity_levels[severity_class]}")
```
For a more complete example, see the inference script in the [GitHub repository](https://github.com/yourusername/home-damage-classifier).
|