A simple Image recognition program using pytorch
This is a simple convolutional neural network capable of recognizing 3 different shapes (house, smiling face and fish)
you can download full version from here
the model is a a small convolutional neural network with two convolution/maxpool and 2 linear layers. the input is a square 500x500 2d array (because all pixels in my datase images are whether (0,0,0) or (255,255,255) it shoul be a 2d array of 0s and 1s). It gives 3 different numbers as output that are then passed through a softmax function to get the probabilities for 3 different shapes
the dataset consists of 243 pictures I drew (81 samples for each shape)
create_dataset.py turns all thesepictures intoa hdf file
here you can download raw images
the model is trained by stochastic gradient descent and mean square error loss function.
here you can download the trained model pth file
first import the conda enviroment
conda env create -f environment.yml
activate the enviroment
conda activate imagerecognition
create a directory named models then download the trained model from here and copy it in models directory now run:
python main.py
the window will open and you can draw. press 'c' on the keyboard to get the results like this:
tensor([-160.3708, 134.9981, 40.6356])
house: tensor(0., grad_fn=<MulBackward0>)
smilingface: tensor(100., grad_fn=<MulBackward0>)
fish: tensor(1.0445e-39, grad_fn=<MulBackward0>)
these are probability percentages for each shape
you can use draw.py script to draw new images for dataset. run:
python draw.py
a window will open. draw a sample then press h,s or f then your sample will be saved in house, smilingface or fish directories respectively