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.ipynb
to train or evaluate- Make sure you have the dataset and models in place
Install required libs with:
pip install torch torchvision matplotlib