Identifying Wagon Numbers Using Transformers

This repository contains the implementation of our Transformer-based approach for identifying railway wagon numbers using raw results from an object detection model. This README provides instructions for repeating training for all models, training the best model architecture, and testing the model.

Repository Structure

README.md: This file containing instructions.
notebooks: Directory containing Jupyter notebooks for exploratory data analysis and model evaluation.
results: Directory where results are stored.
uic_dataset.py: Script for handling the dataset of UIC numbers.
model_stats.py: Script for calculating model statistics.
paper_content: Directory containing the content related to the paper publication.
test_inference.py: Script for running inference tests on the trained model.
datasets: Directory containing the datasets.
models: Directory where trained models are saved.
models.py: Script defining the model architectures.
pyproject.toml: Configuration file for the project.
train.py: Script for training the models.
train_best_arch_long.py: Script for training the best model architecture for an extended period.
requirements.txt: List of dependencies required for the project.

Prerequisites

Ensure you have Python 3.8 or above installed. Install the required dependencies using the following command:

pip install -r requirements.txt

Dataset Preparation

Download the dataset from the provided link in the paper.
Extract the dataset and place it in the data directory.
Ensure the dataset is organized as specified in the uic_dataset.py script.

Training All Models

To train all models, execute the following command:

python train.py

This script will train multiple models defined in models.py and save the trained models in the models directory.

Training the Best Model Architecture

To train the best model architecture for an extended period, use the following command:

python train_best_arch_long.py

This script trains the best-performing model identified during the initial training phase for an extended period to improve performance further.

Testing the Model

To test the trained model, run the following command:

python test_inference.py

This script will load the trained model from the models directory and perform inference on the test dataset. The results will be saved in the results directory.

Additional Information

For more detailed explanations of the model architectures and training process, refer to the notebooks directory which contains Jupyter notebooks used for exploratory data analysis and model evaluation.

For any questions or issues, please refer to the contact information provided in the paper or open an issue in the GitHub repository.

We hope this repository and the instructions provided here help you in reproducing our results and further exploring the potential of Transformer-based models for industrial object detection tasks.

Glaadiss/AssemblyTransformer