Model_maker_training

class Model_maker_training.DataLoader(csv_file_path, images_dir)
load_data()

Loads data from the csv file. :param csv_file_path: The file path to the csv file. :param images_dir: The directory that contains images.

Returns: A Dataloader containing the data of training, validation and test set

class Model_maker_training.ModelTrainer(train_data, validation_data)

A class to train the model on the given dataset and evaluate it.

train_data

Training dataset, in the form of Dataloader.

validation_data

Validation dataset, in the form of Dataloader.

eval

The evaluation result.

model

The trained model.

train_model()

Fine tunes the EfficientDet-Lite0 model on the given dataset.

Parameters:
  • train_data – Training dataset, in the form of Dataloader.

  • validation_data – Validation dataset, in the form of Dataloader.

Returns:

A trained model.

evaluate_and_export(test_data)

Evaluates the model and exports it to the export directory.

Parameters:

test_data – Test dataset, in the form of Dataloader.

class Model_maker_training.ObjectDetector(model, model_path, classes)

A class to perform object detection with a given model.

model

The trained model.

model_path

The file path to the trained model.

classes

A list of class labels.

colors

A list of colors for visualization.

interpreter

The TensorFlow Lite interpreter.

preprocess_image(image_path, input_size)
detect_objects(image, threshold=0.5)
run_odt_and_draw_results(image_path, interpreter, threshold=0.5)

Run object detection on the input image and draw the results.

Parameters:
  • image_path – The file path to the input image.

  • interpreter – The TensorFlow Lite interpreter.

  • threshold – The minimum confidence score for detected objects.

Returns:

A NumPy array of the input image with the detection results.