Search code examples
tensorflowcomputer-visionobject-detectionyolopre-trained-model

Add custom classes to pre-trained data-set


I use the already trained(pre-trained) data-set for object detection using yolo+tensorflow. My inference results are great but now I want to "add" a new class to pre-trained data-set. There are 80 classes in pre-trained data-set how can I add my custom classes and made it 81 or 82 in total? Inference git-hub "https://github.com/thtrieu/darkflow".


Solution

  • In case of transfer learning, pre-trained weights on famous datasets like 'Imagenet', 'fashion-mnist' etc are used. These datasets have defined number of classes and labels which may or may not be same as our dataset. The best practice is to add layers above the output layer of the pre-trained model output. For example in keras:

    from tensorflow.keras.applications import mobilenet
    from tensorflow.keras.layers import Dense, Flatten
    output = mobilenet(include_top=False)
    flatten = Flatten()(output)
    predictions = Dense(number_of_classes, activation='softmax')(layer)
    

    In this case you need to train(or better call it fine tune) the model using your dataset. The mobilenet network will use pretrained weights and the last layer will be only trained as per your dataset with the your defined number of classes.

    You may also use:

    from tensorflow.keras.applications import mobilenet
    preds = mobilenet(include_top=Flase, classes=number_of_classes, weights='imagenet')
    

    for more information you can refer: keras-applications and these blog1, blog2