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 datasetarchive/Dataset/→ Backup copy of the datasetmodels/rock_resnet18_mainclass.pth→ predicts main rock classes (3)rock_resnet18_subclass.pth→ predicts detailed subclasses (7)
notebooks/resnet18.ipynb→ model trainingpreprocessing.ipynb→ resizing, normalization, splittingdata_augmentation.ipynb→ flipping, rotation, brightness
Open the notebooks in Jupyter or VS Code and run:
resnet18.ipynbto train or evaluate- Make sure you have the dataset and models in place
Install required libs with:
pip install torch torchvision matplotlib