Poké Model

A bunch of Python scripts to create a Pokémon Classifier.
Report a bug or request a feature

Table of Contents

• Introduction
• Getting Started
  • Prerequisites and dependencies
  • Installation
• Usage
• Contributing
• License

Introduction

This repository is heavily inspired by the project pokemon-classifier by imjeffhi4, check it out!

The repository contains different scripts to create a Pokémon Classifier. You will find scripts to:

• Collect images of Pokémon on internet to create a dataset.
• Augment the variety of the dataset by doing random transformations to the images.
• Fine-tune a classifier with the dataset.
• Test the classifier.

The final model can be found on Hugging Face.

Getting Started

Prerequisites and dependencies

This repository is tested on Python 3.7+.

You should install Poké Model in a virtual environment. If you're unfamiliar with Python virtual environments, check out the user guide. First, create a virtual environment with the version of Python you're going to use and activate it.

You can install directly all required packages by using the file requirements.txt and doing:

pip install -r requirements.txt

For a more step by step approach, depending of the scripts you want to execute, you need to instal various packages:

Data Collection

For the data collection scripts, the following packages are required :

• selenium is used to do search and retrieve images from the browser Chrome (you need to put the chrome driver executable in the folder DataCollection/chromedriver)
• pillow is used to load and manipulate images.
• cv2 is used to perform transformations on images.
• progress is used to feedback progression with beautiful progress bars.

Training

For the training scripts, the following packages are required :

• pytorch-lightning (and torch) are used to manipulate the model
• transformers is used for the model
• pillow is used to load and manipulate images.
• progress is used to feedback progression with beautiful progress bars.

Testing

For the training scripts, the following packages are required :

• PySide6 is used for the user interface
• torch is used to manipulate the model
• transformers is used for the model
• pillow is used to load and manipulate images.

All packages can be installed with the regular:

pip install <package-name>

Installation

Follow the instructions above then clone the repo (git clone https:://github.com/torresflo/Poke-Model.git).

Usage

To start the data collection, use the script main_collecting.py.
For the training of the model, use the script main_training.py.
Finally, to test the generated model in a simple UI, use the script Testing\main.py.

The images are collected using a mixture of PokeAPI and images scraped from Brave Search.
Note that Brave Search will sometimes requires a Captcha to be performed. The script will wait until it is done.

Training is performed by using PyTorch Lightning and by fine-tuning a ViT-base model (google/vit-base-patch16-224).

Screenshot

Here is a screenshot of the test program (Testing\main.py):

Contributing

Contributions are what make the open source community such an amazing place to be learn, inspire, and create. Any contributions you make are greatly appreciated.

1. Fork the Project
2. Create your Feature Branch (git checkout -b feature/AmazingFeature)
3. Commit your Changes (git commit -m 'Add some AmazingFeature)
4. Push to the Branch (git push origin feature/AmazingFeature)
5. Open a Pull Request

License

Distributed under the GNU General Public License v3.0. See LICENSE for more information.