---
license: bsd-3-clause
library_name: torch
pipeline_tag: "image-classification"
datasets: DanielCruz09/natural_disasters
---

# Natural Disaster Image Classifier

A PyTorch CNN that can classify images into the following categories:
- Fire
- Flood
- Earthquake
- Non-Damage

This model is based on PyTorch's ResNet50 architecture, which is licensed under the BSD-3-Clause License. The original ResNet architecture was introduced in *Deep Residual Learning for Image Recognition* (He et. al., 2015).

# Model Architecture

This model is built on top of **ResNet50**, a deep convolutional neural network. ResNet50 uses stacked residual blocks to enable the training of very deep networks by reducing the effects of vanishing gradients (which are very apparent in several modern architectures) via identity skip connections.

For this project:
- The base architecture comes from **TorchVision's ResNet50** implementation.
- The final fully connected layer was replaced with a classifier tailored to disaster relief image categories.
- All convolutional layers and residual blocks are identical to the original architecture.
- Pretrained ImageNet weights were used for initialization to accelerate convergence and improve generalization.

This architecture provides a robust balance of accuracy and efficiency for real-world image classification tasks.

# Training Details

The model was trained using a custom dataset of natural disaster images, organized into multiple classes as stated above. The training pipeline includes:

## Data Preprocessing
- Images were resized to 224 x 224.
- Normalization followed ImageNet mean & std.

## Training Procedure
- Framework: **PyTorch**
- Optimizer: **SGD**
- Learning Rate: `1e-2`
- Loss function: **Cross-Entropy Loss**
- Batch size: 32

# Evaluation

Model performance metrics computed for this model were:
- Accuracy
- Precision, Recall, F1 (all using a weighted average due to multiclass classification)

# How to Use

## Install the library

```python
pip install huggingface_hub
```

## Download the model

```python
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(
    repo_id="DanielCruz09/disaster-image-classifier",
    filename="model_weights.pth"
)

print("Model downloaded to: ", model_path)
```

## Load the model into PyTorch

```python
import torch
import json
from resnet50 import ResNet50

with open("class_names.json", "r") as f:
    class_names = json.load(f)

mapping = {name: idx for idx, name in enumerate(class_names)}

model = ResNet50(num_classes=len(class_names), mapping=mapping)
state_dict = torch.load(model_path, map_location="cpu")
model.model.load_state_dict(state_dict["model_state_dict"])
```

## Evaluate the model

```python
test_path = "data/processed/Test/"
test_dataset = NaturalDisasterDataset(root=test_path)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=True)

model.eval(test_loader=test_loader, write_path="results/resnet50_results.csv")
```

Alternatively, run **main.py**.

```bash
python -m main.py
```

## Run inference

```python
from torchvision import transforms
from PIL import Image

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        # same as ImageNet
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

img = Image.open("example.jpg").convert("RGB")
image_tensor = preprocess(img).unsqueeze(0)

with torch.no_grad():
    outputs = model(image_tensor)
    probs = torch.softmax(outputs, dims=1)

predicted_class = probs.argmax().item()
print("Predicted class: ", predicted_class)
```

# Reference
He, K., Zhang, X., Ren, S., & Sun, J. (2015). *Deep Residual Learning for Image Recognition.* CVPR. https://arxiv.org/abs/1512.03385