This project is a PyTorch-based Convolutional Neural Network (CNN) for classifying simple geometric shapes (circle, rectangle, square, ellipse, triangle) from grayscale PNG images.
🔁 Looking for the TensorFlow version? Check it out here: CNN Shape Classifier (TensorFlow)
├── dataset.py # Custom PyTorch Dataset for loading shape images
├── labels.py # Shape label mapping
├── model.py # CNN model definition
├── train.py # Model training script
├── test.py # Model evaluation script
├── predict.py # Single image prediction script
├── train/ # Training dataset folder (64x64 PNG images)
├── test/ # Testing dataset folder (64x64 PNG images)
- Python 3.8+
- PyTorch
- torchvision
- Pillow
Install dependencies with:
pip install torch torchvision pillow
- The
train/
andtest/
directories contain the training and testing images. - Images should be 64×64 PNGs, named with the shape name and a number (e.g.,
circle60.png
).
Train the model with:
python train.py
This will:
- Load images from
train/
- Train the CNN (default: 300 epochs — adjust in
train.py
) - Save the trained model as
CNN_model.pth
Evaluate the model’s accuracy on a folder of images:
python test.py
This will:
- Load images from
train/
(you can change the folder intest.py
) - Print predictions and accuracy statistics
Predict the shape in a single image:
python predict.py
Edit the img_path
variable in predict.py
to point to your image.
Example Output
Predicted: triangle (class 4)
See labels.py
for the mapping of shape names to integer labels.
See model.py
for the full CNN definition.