Search code examples
pythontensorflowobject-detection-api

Tensorflow object detection choose/switch class to detect

This question shows how to choose one class to detect. How could you change the code to switch on what classes are displayed? Example filter so only people are displayed then change so only cars are displayed and so on.

Solution

  • Say you have scores and classes:

    import numpy as np

scores = sorted(np.random.rand(100), reverse=True)
classes = np.random.randint(0, 4, 100)

    You can turn the probabilities that are not from class 1 to 0:

    modified_scores = [a if b == 1 else 0 for a, b in zip(scores, classes)]

    It will not display classes that are either 1) below threshold and 2) not from the desired class.