- Generate a synthetic dataset of shapes, such as rectangles, triangles, and ellipses, with random position, orientation, and color. The dataset should be sufficiently large and varied to allow for robust training and testing of the model.
Classification CNN implementing LeNet structure with pytorch
- Step 1 - Install all packages
$ pip install -r requirements.txt
-
Step 2 (OPTIONNAL) - If you wish, you can generate the dataset, running generate_dataset.py will create dataset folder with 3000 images for each shape (rectangles, ellipses and triangles).
-
Step 3 - Make predictions on new shape images (run predict.py will open a window showing a random shape image with the model's prediction). To make a new prediction, tab any keyboard key. To quit the prediction's window, tap 'q'.
The output folder contains the trained model (model.pth) and its performance (plot.png)
Here is the screenshot result of the prediction that you should have when you run predict.py.
