Code for importing mutiny dataset into Voxel51

mutiny.py

import fiftyone as fo
import os
import xml.etree.ElementTree as ET
import cv2

name = "headsegmentation_dataset_ccncsa"
dataset_dir = "/home/priyansh-kedia/Documents/Datasets/headsegmentation_dataset_ccncsa/"

dataset = fo.Dataset()

sub_dirs = os.listdir(dataset_dir)

tree = ET.parse(os.path.join(dataset_dir, 'training.xml'))

root = tree.getroot()
mask_path = ""
image_path = ""
    
def create_detections(element, image_path):
    dets = []
    if element.tag == 'bboxes':
        for e in element.iter("box"):
            label = e.attrib["label"]

            xmin = float(e.attrib["left"])
            ymin = float(e.attrib["top"])
            xmax = float(xmin + float(e.attrib["width"]))
            ymax = float(ymin + float(e.attrib["height"]))
            
            metadata = fo.ImageMetadata.build_for(image_path)

            w = metadata.width
            h = metadata.height

            bbox = [xmin / w, ymin / h, (xmax - xmin) / w, (ymax - ymin) / h]
            detection = fo.Detection(bounding_box=bbox, label=label)
            dets.append(detection)
    detections = fo.Detections(detections=dets)
    return detections

def create_segmentations(element, mask_path, image_path):
    segs = []
    if element.tag == 'bboxes':
        mask = None
        for e in element.iter("box"):
            label = e.attrib["label"]

            if label in ["1","2","3","4","6","7"]:
                
                xmin = float(e.attrib["left"])
                ymin = float(e.attrib["top"])
                xmax = float(xmin + float(e.attrib["width"]))
                ymax = float(ymin + float(e.attrib["height"]))

                if label in ["6","2","3"]:
                    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) < 122
                if label in ["1"]:
                    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) > 140 
                if label in ["4"]:
                    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) < 179
                if label in ["7"]:
                    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
                
                metadata = fo.ImageMetadata.build_for(image_path)
                
                w = metadata.width
                h = metadata.height
            
                bbox = [xmin / w, ymin / h, (xmax - xmin) / w, (ymax - ymin) / h]
                cropped_mask = mask[int(ymin):int(ymax), int(xmin):int(xmax)]
                segmentation = fo.Detection(bounding_box=bbox, label=label, mask=cropped_mask)
                segs.append(segmentation)
    segmentations = fo.Detections(detections=segs)
    return segmentations



for element in root.iter():
    if element.tag == 'srcimg':
        image_path = os.path.join(dataset_dir, element.attrib['name'])
        image_path = image_path.replace("\\","/")
    if element.tag == 'labelimg':
        mask_path = os.path.join(dataset_dir, element.attrib['name'])
        mask_path = mask_path.replace("\\","/")
    if image_path and mask_path:
        sample = fo.Sample(image_path)

    if element.tag == 'bboxes':
        detections = create_detections(element, image_path)
        sample["detections"] = detections
        segmentations = create_segmentations(element, mask_path, image_path)
        sample["segmentations"] = segmentations
        dataset.add_sample(sample)

session = fo.launch_app(dataset, desktop=False)
session.wait()