Skip to content

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

🪨 Rock Classification with ResNet18

A simple PyTorch project that classifies rocks by image using a ResNet18 model. The model can predict both:

  • Main class: Igneous, Metamorphic, Sedimentary
  • Subclass: Basalt, Granite, Marble, Quartzite, Coal, Limestone, Sandstone

📁 Folder Structure

  • Dataset/ → Rock image dataset
  • archive/Dataset/ → Backup copy of the dataset
  • models/
    • rock_resnet18_mainclass.pth → predicts main rock classes (3)
    • rock_resnet18_subclass.pth → predicts detailed subclasses (7)
  • notebooks/
    • resnet18.ipynb → model training
    • preprocessing.ipynb → resizing, normalization, splitting
    • data_augmentation.ipynb → flipping, rotation, brightness

🧪 How to Use

Open the notebooks in Jupyter or VS Code and run:

  • resnet18.ipynb to train or evaluate
  • Make sure you have the dataset and models in place

🔧 Requirements

Install required libs with:

pip install torch torchvision matplotlib

About

This project uses a custom-trained ResNet18 to classify different types of rocks into 7 subclasses and optionally into 3 main rock types: Igneous, Metamorphic and Sedimentary. It's built with PyTorch and tested on a dataset of rock images.

Resources

Stars

0 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages