how to save the detectron2 prediction label result?

[Suzansb]

I am working detectron2 object detection model, which produced a good output result. I can also see the output predicted result with labels with the help of detectron2.utils.visualizer Visualizer functions like the following:

But I want to print the predicted labels and save them in txt or csv file in the direction from top left side. Ex. Dhaka Metro Ga 128691 How can I do this?

[daksh47]

its been 2-3 years since the question was asked, but still going to answer it as it might provide an reference to others

• first we iterate the predictions to get the coords of all the bbox's ( order doesnt matter )
• while iterating we calculate the center points or center coords of each bbox by centerX = (top-right x value + bottom-left x value) / 2 and centerY = (top-right y value + bottom-left y value) / 2
• [ 100, 50, 150, 200 ] => this is how the output of each bbox coords will be, the first two values (x1,y1) = (100,50) are either top-right OR top-left values ( it doesnt matter for our use case, as we are finding the mid point), and the last two values (x2,y2) = (150,200) are either bottom-right OR bottom-left values
• after calculating the midpoint of each bbox using dictionary data structure (key:value pairs) we store midpoint of each bbox and its corresponding predicted class
• then we find the height of the image on which the predictions were done
• using the height/2 value as a medium to classify the predictions as upper or lower, we create two sorted lists each having upper and lower coords respectively which are later concatenated, upper and lower coords is divided by comparing the Y coord value with the height value and the sorting of the divided lists are done by comparing the X coord value, leading to a perfect arrangement of coords in mentioned order by you
• finally having a single list of coords in your required order, that is top-left -> top-right -> bottom-left -> bottom-right, we iterate over the list and get corresponding class index, then we use the class index as an index in the list of classes used to train the model, this is done to get the corresponding class name

in simple words, this problem is solved by dividing the image into two halves upper and lower, and the bbox of upper and lower halves are sequentially arranged from left to right based upon their position values in each halves

pointsAndClass = dict()  # Dict to store coords and corresponding class
for i in range(len(predictions["instances"])):  # iterate over all the predictions
    pred_box = predictions["instances"].pred_boxes[i].tensor.tolist()[0]  # Get bounding box coordinates
    pred_class_index = predictions["instances"].pred_classes[i].item()  # Get corresponding class index for bbox
    pred_box = [int(x) for x in pred_box]  # Convert float coords to int

    centerPointX = int((pred_box[0] + pred_box[2]) / 2)  # get the x value of center point of bbox
    centerPointY = int((pred_box[1] + pred_box[3]) / 2)  # get the x value of center point of bbox

    pointsAndClass[(centerPointX, centerPointY)] = pred_class_index  # add the center coords of each bbox along with their corresponding class index

height, width, channels = image.shape  # get the dimensions of the images
value = height / 2  # to create comparison value to classify the coords into upper or lower part of the image

coords = list(pointsAndClass.keys())  # get the list of center coords of each bbox
orderdCoords = sorted(list([(x, y) for x, y in coords if y < value]), key=lambda x: x[0]) + sorted(list([(x, y) for x, y in coords if y >= value]),key=lambda x: x[0])  # two sorted lists have been created based upon the comparison value and later concatenated

class_names = MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]).thing_classes  # replace with your dataset ( which includes your classes example ["cat", "dog"], which is basically an array containing all your classes )
for i in orderdCoords:  # iterate over the classified and sorted coords
    print(class_names[pointsAndClass[i]], end=" ")  # returns the corresponding class name for the class index

this might or might not be the best way to solve it, have just mentioned an workaround which is accurate for your requirement and few other cases too, for example when viewed this looks like a model to detect number plates in the country Bangladesh, various real world, such as the whole predictions in a single half, or the format of the number or the codes varying from vehicle to vehicle, all these cases have been considered, to be sure and confident make sure that the predicting image supplied as input just contains the number plate to avoid false positives

ALSO FOR THE PRECISE WORKING OF THIS TECHNIQUE MAKE SURE THAT THE DISTANCE FROM THE BBOX OF THE UPPER HALF TO THE TOP OF THE IMAGE IS THE SAME OR SIMILAR TO THE DISTANCE OF THE BBOX OF THE BOTTOM HALF TO THE BOTTOM OF THE IMAGE

for example:

-------- top of image
|
| 10cm or 100px or whatever you find comfortable
|
bbox bbox

bbox bbox
|
| 10cm or 100px or whatever you find comfortable
|
---------- bottom of image

Any questions related to code, logic, Code not working for few extreme cases etc.. feel free to contact