Skip to content

Built a neural network from scratch and trained it to recognize geometric shapes (circle, ellipse, square, rectangle, triangle) drawn on a 64x64 grid using PyTorch.

Notifications You must be signed in to change notification settings

JacemHaggui/cnn-shape-classifier-pytorch

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

CNN Shape Classifier

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)

Project Structure

├── 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) 

Requirements

Install dependencies with:

pip install torch torchvision pillow

Dataset

  • The train/ and test/ 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).

Training

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

Testing

Evaluate the model’s accuracy on a folder of images:

python test.py

This will:

  • Load images from train/ (you can change the folder in test.py)
  • Print predictions and accuracy statistics

Predicting a Single Image

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)

Label Mapping

See labels.py for the mapping of shape names to integer labels.

Model Architecture

See model.py for the full CNN definition.

About

Built a neural network from scratch and trained it to recognize geometric shapes (circle, ellipse, square, rectangle, triangle) drawn on a 64x64 grid using PyTorch.

Topics

Resources

Stars

Watchers

Forks

Contributors 2

  •  
  •  

Languages