Skip to content

Implementation of a neural network for solving the problem of mushroom classification based on the TensorFlow library and the pre-trained efficientnet/b0 model.

License

Notifications You must be signed in to change notification settings

ruddanil/NN-MushroomClassification-Py-TensorFlow-SSTU

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

8 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

NN-MushroomClassification-Py-TensorFlow-SSTU

Implementation of a neural network for solving the problem of mushroom сlassification based on the TensorFlow library and the pre-trained efficientnet/b0 model. The Tkinter framework is used to draw the interface.

image

Using the program to solve other сlassification tasks

Option 1. You want to train/retrain the model from scratch on your dataset.

The program implements the neural network training functionality on any image dataset. The only condition is that the data in the dataset must be divided into folders containing images of a certain category.

To train on your data, you need to do the following:

  1. In the source code of the program, replace the category headers with those contained in your dataset.
dataset_labels = ['Альбатреллус Овечий', 'Аурикулярия', 'Белый Гриб', 'Боровик', 'Вешенка',
                  'Волнушка', 'Гигрофор', 'Головач', 'Груздь', 'Зимний Гриб', 'Зонтик',
                  'Лиофиллум Ильмовый', 'Лисичка', 'Масленок', 'Мокруха', 'Моховик',
                  'Подберезовик', 'Подосиновик', 'Рыжик', 'Спарассис Курчавый']

image

  1. Clear or copy the contents of the "saved_models" folder to another directory.
  2. Click on the "Dataset" button and select the directory where the dataset is located.
  3. Specify the number of training epochs.
  4. Press the "Train" button. The trained model will be saved in the saved_models folder.
  5. You can start the recognition. Click the "Choose Image" button, select your image and click the "Recognize Image" button.

As a result, the top 3 most likely recognized categories are displayed at the bottom of the screen, indicating confidence in %.

Optional: You can replace the pre-trained TensorFlow model by specifying a link to it in the source code.

    model = tf.keras.Sequential([
        hub.KerasLayer("https://tfhub.dev/tensorflow/efficientnet/b0/feature-vector/1",
                       output_shape=[1280],
                       trainable=False),
        tf.keras.layers.Dropout(0.4),
        tf.keras.layers.Dense(train_generator.num_classes, activation='softmax')
    ])

Note: This program is an educational project. It is convenient to use it to demonstrate the work of a pre-trained model. To solve more serious problems, it is more convenient and correct to use tools like "Jupiter Notebook", where the learning progress is displayed and additional settings are possible.

Option 2. You have a pre-trained model.

  1. Replace dataset_labels in the source code (by analogy with option 1.1).
  2. Click the "Load Model" button and specify the directory with your pre-trained model.
  3. Proceed with recognition.

About

Implementation of a neural network for solving the problem of mushroom classification based on the TensorFlow library and the pre-trained efficientnet/b0 model.

Topics

Resources

License

Stars

Watchers

Forks