Search code examples
computer-visionobject-detectionyolo

Building an object detector for a small dataset with a single class


I have a dataset of a single class (rectangular object) with a size of 130 images. My goal is to detect the object & draw a circle/dot/mark in the centre of the object.

Because the objects are rectangular, my idea is to get the dimensions of the predicted bounding box and take the circle/dot/mark as (width/2, height/2).

However, if I were to do transfer learning, would YOLO be a good choice to detect a single class of objects in a small dataset?


Solution

  • YOLO should be fine. However it is old now. Try YoloV4 for better results. People have tried transfer learning from FasterRCNN to detect single objects with 300 images and it worked fine. (Link). However 130 images is a bit smaller. Try augmenting images - flipping, rotating etc if you get inferior results.

    Use same augmentation for annotation as well while doing translation, rotation, flip augmentations. For example in pytorch, for segmentation, I use:

        if random.random()<0.5:   # Horizontal Flip
            image = T.functional.hflip(image)
            mask  = T.functional.hflip(mask)
    
        if random.random()<0.25:  # Rotation
            rotation_angle = random.randrange(-10,11)
            image = T.functional.rotate(image,angle = rotation_angle)
            mask  = T.functional.rotate(mask ,angle = rotation_angle)
    

    For bounding box you will have to create coordinates, x becomes width-x for horizontal flip.

    Augmentations where object position is not changing: do not change annotations e.g.: gamma intensity transformation