igorastashov / mlops-example-project Goto Github PK

Pokemon image classification using Mobilenet_v2 model

Astashov I.V., 2023.

This repository contains model, evaluation code and training code on dataset from kaggle. If you would like to run pretrained model on your image see (2) Quick start.

• This template use poetry to manage dependencies of your project. First you need to install poetry;
• Next run poetry install and poetry shell.

# Download model and optimizer
cd weights
bash download_weights.sh
cd ../..

This command runs the model on a single image, and outputs the prediction. Put your Pokémon image into the appropriate folder photo to test.

# Model weights need to be downloaded
python inference.py

A script for downloading the dataset is as follows:

# Download the dataset
cd data
bash download_data.sh
cd ../..

All the 150 Pokémon included in this dataset are from generation one. There are around 25 - 50 images for each Pokémon. All of them with the Pokémon in the center. Most (if not all) of the images have relatively high quality (correct labels, centered). The images don't have extremely high resolutions so it's perfect for some light deep learning.

If the script doesn't work, an alternative will be to download the zip files manually from the link. One can place the dataset zip files in data, respectively, and then unzip the zip file to set everything up.

PAY ATTENTION

This repository runs Data Version Control (DVC) for training and validation data. Pre-configured Google Drive remote storage stores raw input data.

$ dvc remote list
my_remote gdrive://1RXz3Mv7OxVveHtQ7c1ZtGgazDh6bPFJz

You can run dvc pull to download the data:

$ dvc pull

Example script to train and evaluate model.

# Train ConvNet
python main.py

Hydra - hyperparameter management tool.

If you want to change hyper-parameters such as: epoch_count, lr, batch_size, momentum, then you can change them in the file conf/config.yaml.

This repository borrows partially from Isadrtdinov, and FUlyankin repositories. Repository design taken from v-goncharenko, PeterWang512 and ArjanCodes.