Search code examples
tensorflowkerastensorflow2.0object-detectionfaster-rcnn

How to use FasterRCNN Openimages v4?


I can't seem to find any documentation on how to use this model. I am trying to use it to print out the objects that appear in a video any help would be greatly appreciated I am just starting out so go easy on me


Solution

  • After some digging around and a LOT of trial and error I came up with this

    #!/home/ahmed/anaconda3/envs/TensorFlow/bin/python3.8
    import tensorflow as tf
    import tensorflow_hub as hub
    import time,imageio,sys,pickle
    
    # sys.argv[1] is used for taking the video path from the terminal
    video = sys.argv[1]
    #passing the video file to ImageIO to be read later in form of frames
    video = imageio.get_reader(video)
    dictionary = {}
    #download and extract the model( faster_rcnn/openimages_v4/inception_resnet_v2 or
    # openimages_v4/ssd/mobilenet_v2) in the same folder
    module_handle = "*Path to the model folder*"
    detector = hub.load(module_handle).signatures['default']
    #looping over every frame in the video
    for index, frames in enumerate(video):
        # converting the images ( video frames ) to tf.float32 which is the only acceptable input format
        image = tf.image.convert_image_dtype(frames, tf.float32)[tf.newaxis]
        # passing the converted image to the model
        detector_output = detector(image)
        class_names = detector_output["detection_class_entities"]
        scores = detector_output["detection_scores"]
        # in case there are multiple objects in the frame
        for i in range(len(scores)):
            if scores[i] > 0.3:
                #converting form bytes to string
                object = class_names[i].numpy().decode("ascii")
                #adding the objects that appear in the frames in a dictionary and their frame numbers
                if object not in dictionary:
                    dictionary[object] = [index]
                else:
                    dictionary[object].append(index)
    print(dictionary)