Upload damage classification model

README.md

@@ -1,3 +1,18 @@
+---
+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
@@ -36,30 +51,61 @@ This model was trained to classify damage to household items, identifying both t
 - 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, ViTForImageClassification
+from transformers import ViTFeatureExtractor
 from PIL import Image
+import torch
 
 # Load model and feature extractor
-model = ViTForImageClassification.from_pretrained("USER/REPO_NAME")
+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 prediction
+# Get predictions
 outputs = model(**inputs)
-predicted_class_idx = outputs.logits.argmax(-1).item()
-
-# Get class name if available
-if hasattr(model.config, "id2label"):
-    predicted_class = model.config.id2label[predicted_class_idx]
-    print(f"Predicted class: {predicted_class}")
-else:
-    print(f"Predicted class index: {predicted_class_idx}")
+
+# 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).