I'm sharing my code snapshot that saves the inference output of Detection2 and exports the prediction polygons to a CSV file. Modifying the last code block can easily allow you to save the polygons to JSON or a shapefile.
The first code block handles inference and saves the predictions (outputs) as PKL files, while also saving the Visualizer plots as PNG images.
import pandas as pd
import glob
test_folder = Path(os.getcwd() + r"\dataset\test") #We need test_folder to be Path object, not string, for test_folder.glob to work
print ('The folder of images: ', test_folder)
for i, imageName in enumerate(test_folder.glob("*.png")):
file = str(imageName)
file_name = file.split("\\")[-1] # split and get the file name only
im = cv2.imread(file)
outputs = predictor(im) #Inference
# attributes variables
predictions_output = {}
predictions_output["file_name"] = file_name
predictions_output["file_location"] = file
predictions_output["prediction"] = outputs
# save the results with iteration no as pickle objects.
predictions_pkl = 'predictions_' + str(i) + '.pkl'
predictions_pkl_file = os.getcwd() + r"\dataset\predictions" + '\\' + predictions_pkl
#print ('## The predictions_pkl_file: ', predictions_pkl_file)
with open(predictions_pkl_file , "wb") as f: # wb opens the file in binary format for writing
pickle.dump(predictions_output, f)
v = Visualizer(
im[:, :, ::-1],
metadata=metadata,
scale=1,
)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
plt.figure(figsize=(20, 20))
plt.imshow(out.get_image()[:, :, ::-1])
# save the visualizer result
saved_image = os.getcwd() + r"\dataset\output_images" + '\\' + file_name
plt.savefig(saved_image)
The 'outputs' variable contains the inference output, including ["pred_boxes"] and ["pred_masks"], in the format described in the Detectron2 tutorial: https://detectron2.readthedocs.io/tutorials/models.html#model-output-format.
The 'pred_masks' values are in binary (False, True). The OpenCV function 'findContours' is used to extract the contours of each predicted mask. Thus, the second code block is responsible for extracting 'pred_boxes,' 'pred_masks,' and contours from the 'outputs' or predictions.
from shapely.geometry import Polygon
csv_df = pd.DataFrame()
for prediction in glob.glob(os.getcwd() + r"\dataset\predictions\*pkl"):
fileName_absolute = os.path.basename(prediction) ## get the file name
print("The pkl_file name: ", fileName_absolute)
# Read the pkl file
pickleFile = open(prediction,"rb")
pickel_Info = pickle.load(pickleFile)
# Get the 'pred_boxes' values in the pkl file:
boxes = (
pickel_Info["prediction"]["instances"]
.get_fields()["pred_boxes"]
.tensor.cpu()
.numpy()
)
# Get the 'pred_mask' values in the pkl file:
masks = (
pickel_Info["prediction"]["instances"]
.get_fields()["pred_masks"]
.cpu()
.numpy()
.astype('uint8')
)
# Extract the 'contour' of each 'pred_box' by converting pred_box boolean to a 8-bit numpy array:
# Refer to this post to get and visualise the contours: https://stackoverflow.com/questions/73217530/extract-the-masks-values-from-detectron2-object-detection-segmentation-and-then
contours = []
for pred_mask in pickel_Info["prediction"]['instances'].pred_masks:
mask = pred_mask.cpu().numpy().astype('uint8')
contour, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
contours.append(contour[0]) # contour is a tuple (OpenCV 4.5.2), so take the first element which is the array of contour points
# save boxes and masks-polygons
dataset = []
counter_box = 0
for i, box in enumerate(boxes):
data = {}
data["file_name"] = pickel_Info['file_name']
data["file_location"] = pickel_Info["file_location"]
data["image_height"] = pickel_Info["prediction"]["instances"].image_size[0]
data["image_width"] = pickel_Info["prediction"]["instances"].image_size[1]
data["bounding_box"] = box
data["mask"] = masks[i]
data["contours"] = contours[i]
data["polygon"] = Polygon(np.squeeze(contours[i]))
data["id"] = "id_{counter_box}"
counter_box = counter_box + 1
dataset.append(data)
df = pd.DataFrame(dataset)
# keep specific columns for the csv file
df = df[
[
"file_name",
"id",
"bounding_box",
"polygon",
]
]
csv_df = csv_df.append(df)
print ('## csv_df: ', csv_df)
The third block involves saving the DataFrame to a CSV file.
csv_df.to_csv(os.getcwd() + r"\dataset\predictions_csv\result_polygons.csv')
I hope this helps you!
