This repository contains tools to train your own model for detecting abnormal human behaviour using the mediapipe framework. The models are built based on autoencoders with different architectures:
- A simple Neural Network with fully-connected layers
- Based on 2D Convolutional Neural Network
- Based on 3D Convolutional Neural Network (if it needs to realize)
- Transformer Based
At the current moment, the mediapipe framework doest not work with the latest version of Python. Repository was developed on Python 3.8, so it's most recommended.
List of using libraries:
- mediapipe 0.10.0
- opencv-python 4.7.0.72
- scikit-learn 1.2.2
- torch 2.0.1
- torchvision 0.15.2
- tqdm 4.65.0
- numpy 1.24.3
The main class in this code that you will need to work with is the Trainer
class. It allows you to specify model type, model parameters, training parameters, add data for training and validation, as well as create a set of abnormal data, and of course train the model.
Initialization of the class looks like this:
from trainer import Trainer
from parameters import init_transformer_params
model = init_transformer_params(init_transformer_params())
The init_transformer_params
function returns a dictionary with the default parameters for the Transformer model. You can change them as you like, or create your own function to initialize the parameters.
You can also specify the following parameters during initialization:
params
- dictionary with parameters for the model. Fromparameters
module you can import functions to initialize parameters for different models.model
: str - model type. Can be one of the following:fc_cnn
,cnn
,transformer
. Default:transformer
.criterin
- loss function. Default:MSELoss()
.learning_rate
- learning rate. Default:0.001
.optimizer
- optimizer. Default:Adam
.device
- device for training. Default:cpu
.batch_size
- batch size. Default:32
.num_epochs
- number of epochs. Default:10
.sequence_length
- length of the sequence. Default:12
.
Note that when you initialize parameters from parameters
, you can also specify different values for model parameters.
For adding data you can use several methods: add_data
, add_validation_data
, add_abnormal_data
, add_anomaly_data
and create_validation_set
. Three first methods are similar in using. You need specify directory from you want adding data and optionality specify file_format
(in default it is mp4
). Here the example of adding data for model above:
model.add_data("path/to/normal/data")
model.add_validation_data("path/to/normal/validation/data")
model.add_anomaly_data("path/to/abnormal/data")
This code will create 3 data sets for the model. It is mandatory to add training data to train the model, the rest if desired. Note that when the data is first initialized, there will be an AU extraction using mediapipe and this will take some time. However, after processing the video, the pkl
files will be created in the specified directory containing the data already required for further work with these videos without reapplying to mediapipe.
Last method create validation set from train_data. It is necessary for validation model during training. You can specify validation_size
(in default it is 0.2
).
model.create_validation_set(validation_size=0.1)
This method allow you choose path and name with which the model will be saved (in default it is default_model
).
model.set_model_name("path/to/model/name")
After you adding all necessary data, you can start training the model. For this you need to call the train
method. You can specify save_model
(in default it is False
) if you want save model in every epoch.
model.train(save_model=True, transpose=(0, 3, 1, 2))
Parameter transpose
you can specify if you work with CNN type model. Note that if you change this parameter from default (0, 3, 1, 2), you also need to initialize parameter with specifying values.
Trainer can memorize best model and save it. For this you need to call the save_best_model
method.
model.save_best_model()
In test
module you can call test
function.
from test import test
test(source="path/to/your/video.mp4",
model_path="path/to/your/model.pt",
params_path="path/to/your/model/params.json",
model_type="transformer",
criterion="MSELoss",
transpose=(0, 3, 1, 2)
)
In source
you can also specify 0
and then the model will take data from your webcam.