i train my model on MASK2FORMER instance Segmentation Lane Marking Detection. now my model give me prediction using checkpoint and model config file. i want to convert those mask into polygons or i want to Create JSON for Each Predicted image and JSON save Separately. I tried to Save JSON but it is not saving JSON file but model working accurately. MODEL: https://github.com/open-mmlab/mmdetection Instance Segmentation: https://github.com/open-mmlab/mmdetection/tree/master/configs/mask2former "Swin-S"
This code is a script that uses the mmdetection library, a popular library for object detection in computer vision, to perform object detection on images, and then saves the detected objects' bounding boxes and labels to a JSON file.
# Copyright (c) OpenMMLab. All rights reserved.
import asyncio
import os
from tqdm import tqdm
import mmcv
import numpy as np
import torch
from shapely.geometry import Polygon
import csv
import pandas as pd
import cv2
from mmdet.apis import (async_inference_detector, inference_detector, init_detector, show_result_pyplot)
import json
**The script first sets several variables at the top such as the source of the images to be processed, the configuration file and checkpoint for the model, and the output directory where the results will be saved. **
SOURCE = "demo/images"
CONFIG = "configs/mask2former_custom/mask2former_swin-s-p4-w7-224_lsj_8x2_50e_coco.py"
CHECKPOINT = "weights/iter_converted_216000.pth"
OUTPUT = "exp/testing"
os.makedirs(OUTPUT, exist_ok=True)
It then creates the output directory if it does not exist. It also sets the device to use for inference, the palette for visualization, the score threshold for bounding box, and whether to use async inference.
DEVICE = "cuda:0" # Device used for inference
PALETTE = "coco" # ['coco', 'voc', 'citys', 'random']
SCORE_THRESHOLD = 0.3 # bbox score threshold
ASYNC_TEST = False # 'whether to set async options for async inference.
The script then defines a function processResutls which takes in the result of the inference and the classes of the objects that the model is trained to detect. The function first separates the bounding boxes and mask results from the input result tuple. It then converts the masks to polygons and draw segmentation masks. It creates a dictionary to store the results, which includes the image name, bounding boxes and labels. It then saves the results to a JSON file.
# def processResutls(result, classes):
def processResutls(result, classes, image_name):
assert isinstance(result, tuple)
bbox_result, mask_result = result
bboxes = np.vstack(bbox_result)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
if len(labels) == 0:
bboxes = np.zeros([0, 5])
masks = np.zeros([0, 0, 0])
# Convert masks to polygons
polygons = []
for i, mask in enumerate(masks):
# Find contours of the mask
contours = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
contours = contours[0] if len(contours) == 2 else contours[1]
# Convert contours to polygons
for contour in contours:
# Flatten the contour and create a shapely polygon object
contour = contour.flatten().tolist()
polygon = Polygon(contour)
polygons.append(polygon)
# draw segmentation masks
else:
masks = mmcv.concat_list(mask_result)
if isinstance(masks[0], torch.Tensor):
masks = torch.stack(masks, dim=0).detach().cpu().numpy()
else:
masks = np.stack(masks, axis=0)
# dummy bboxes
if bboxes[:, :4].sum() == 0:
num_masks = len(bboxes)
x_any = masks.any(axis=1)
y_any = masks.any(axis=2)
for idx in range(num_masks):
x = np.where(x_any[idx, :])[0]
y = np.where(y_any[idx, :])[0]
if len(x) > 0 and len(y) > 0:
bboxes[idx, :4] = np.array(
[x[0], y[0], x[-1] + 1, y[-1] + 1],
dtype=np.float32)
# create a dictionary to store the results
results = {'image_name': image_name, 'bboxes': bboxes.tolist(), 'labels': labels.tolist()}
# save the results to a JSON file
with open(f'{OUTPUT}/{image_name}.json', 'w') as f:
json.dump(results, f)
assert bboxes is None or bboxes.ndim == 2, \
f' bboxes ndim should be 2, but its ndim is {bboxes.ndim}.'
assert labels.ndim == 1, \
f' labels ndim should be 1, but its ndim is {labels.ndim}.'
assert bboxes is None or bboxes.shape[1] == 4 or bboxes.shape[1] == 5, \
f' bboxes.shape[1] should be 4 or 5, but its {bboxes.shape[1]}.'
assert bboxes is None or bboxes.shape[0] <= labels.shape[0], \
'labels.shape[0] should not be less than bboxes.shape[0].'
assert masks is None or masks.shape[0] == labels.shape[0], \
'segms.shape[0] and labels.shape[0] should have the same length.'
assert masks is not None or bboxes is not None, \
'segms and bboxes should not be None at the same time.'
if SCORE_THRESHOLD > 0:
assert bboxes is not None and bboxes.shape[1] == 5
scores = bboxes[:, -1]
inds = scores > SCORE_THRESHOLD
bboxes = bboxes[inds, :]
labels = labels[inds]
if masks is not None:
masks = masks[inds, ...]
labels = [classes[i] for i in labels]
bboxes = np.array(bboxes)
scores = bboxes[:, -1]
bboxes = bboxes[:, :4].astype(np.int32)
return [scores, labels, bboxes, masks]
def main():
# build the model from a config file and a checkpoint file
model = init_detector(CONFIG, CHECKPOINT, device=DEVICE)
config = mmcv.Config.fromfile(CONFIG)
for image_name in tqdm(os.listdir(SOURCE)):
image_path = os.path.join(SOURCE, image_name)
output_path = os.path.join(OUTPUT, image_name)
# test a single image
result = inference_detector(model, image_path)
# scores, labels, bboxes, masks = processResutls(result, config.classes)
# show the results
show_result_pyplot(
model,
image_path,
result,
palette=PALETTE,
score_thr=SCORE_THRESHOLD,
out_file=output_path
)
async def async_main():
# build the model from a config file and a checkpoint file
model = init_detector(CONFIG, CHECKPOINT, device=DEVICE)
# test a single image
tasks = asyncio.create_task(async_inference_detector(model, SOURCE))
result = await asyncio.gather(tasks)
# show the results
show_result_pyplot(
model,
SOURCE,
result[0],
palette=PALETTE,
score_thr=SCORE_THRESHOLD,
out_file=OUTPUT)
if __name__ == '__main__':
if ASYNC_TEST:
asyncio.run(async_main())
else:
main()`