Search code examples
pythonpytorchfiftyone

FiftyOne Save Image with Overlaid Detections


I am using Pytorch and Fiftyone to process image detections and then visualize these image detections around people like so:

enter image description here

However, I am having difficulty saving this in an easily viewable manner. I want to be able to save the processed image with the bounding boxes overlaid onto the image through the script, which I can only do now by right clicking and downloading the image from the application above. FiftyOne provides multiple options for exporting data: https://voxel51.com/docs/fiftyone/user_guide/export_datasets.html#supported-formats, but all of these export the detection for use in another script (by saving the images and detections seperately in a .txt/.json/etc file) rather than a 'final visualization' image. How can I save the image you see above (including the detection boxes) using FiftyOne? If there is no built in method, can I export it to another type of dataset and save the detections there?


Solution

  • FiftyOne has this capability built-in allowing you to draw labels on samples and save them to disk for any dataset, view, or even just individual samples: https://voxel51.com/docs/fiftyone/user_guide/draw_labels.html

    In general, it can be as simple as:

    import fiftyone as fo
    
    # The Dataset or DatasetView containing the samples you wish to draw
    dataset_or_view = fo.Dataset(...)
    
    # The directory to which to write the annotated media
    output_dir = "/path/for/output"
    
    # The list of `Label` fields containing the labels that you wish to render on
    # the source media (e.g., classifications or detections)
    label_fields = ["ground_truth", "predictions"]  # for example
    
    # Render the labels!
    dataset_or_view.draw_labels(output_dir, label_fields=label_fields)
    

    The draw_labels() method also accepts a DrawConfig that provides a lot of options for how to render the labels when drawing them.