DanielCruz09/disaster-image-classifier

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

pip install huggingface_hub

Download the model

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

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

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.

python -m main.py

Run inference

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