This repository contains the code for training a classification/retrieval model to identify anime titles from scene images.
The goal of this project is to address the issue of identifying anime titles from images posted on the internet. This system enables users to upload an image of an anime scene and receive the corresponding anime title.
• Classification Model: Classifies anime based on scene images.
• Retrieval Model: Retrieves the most similar anime titles for given scene images.
Because the dataset file size is huge, I managed to put it on Google Drive.
Dataset drive : https://drive.google.com/file/d/1adYN9lmDqZ0vXpLLiJJ1kvilKKUazvdw/view?usp=sharing
This repository includes implementations of several popular neural network architectures for efficient image classification and retrieval:
• EfficientNet-B0
• EfficientNet-B7
• ResNeXt-101
• SEResNeXt-101
• MobileNetV2
A custom label smoothing loss function is implemented to improve model generalization:
• Label Smoothing Loss
The main.py script includes the training and validation pipeline, utilizing data loaders, augmentation techniques, and model checkpoints to facilitate efficient training.
The feature_extraction.py script is designed to extract deep features from the trained models, which can then be used for image retrieval tasks.
The retrieve.py script demonstrates how to perform image retrieval using extracted features and a query image. It displays the most similar images from the dataset.
The classification_evaluation.py and retrieval_evaluation.py scripts provide a comprehensive evaluation of the models using the following metrics:
• Accuracy
• Top-5 Accuracy
• Mean Reciprocal Rank (MRR)
• Mean Average Precision (MAP)
Requirements
• Python 3.x
• PyTorch
• torchvision
• numpy
• scikit-learn
• PIL
• tqdm
- Clone the repository:
git clone https://github.com/yourusername/your-repo-name.git
cd your-repo-name
- install the required dependencies:
pip install -r requirements.txt
To start training a model, run:
python main.py --train
To extract features from a trained model, run:
python feature_extraction.py
To perform image retrieval, run:
python retrieve.py --query <path_to_query_image>
To evaluate the model performance on classification tasks, run:
python classification_evaluation.py
To evaluate the model performance on retrieval tasks, run:
python retrieval_evaluation.py
For detailed explanations and insights about this project, visit my medium blog(Thai) : https://medium.com/@thanaphum04571/anime-recognition-system-from-scene-images-anirec-ระบุชื่อเรื่องจากภาพซีนอนิเมะ-c0133dbbfaa9