Search code examples
pythondictionarydata-structuresroboflowinstance-segmentation

How to restructure instance segmentation predictions into a custom dictionary format in Python?


I'm performing instance segmentation using a model trained in RoboFlow, for the prediction result I'm getting:

[InstanceSegmentationInferenceResponse(visualization=None, frame_id=None, time=None, image=InferenceResponseImage(width=720, height=1280), predictions=[
InstanceSegmentationPrediction(x=352.0, y=569.0, width=420.0, height=1066.0, confidence=0.9057266712188721, class_name='animal', class_confidence=None, points=[Point(x=244.125, y=36.0), Point(x=243.0, y=38.0), Point(x=241.875, y=38.0), Point(x=237.375, y=46.0), Point(x=236.25, y=46.0), Point(x=232.875, y=52.0), Point(x=232.875, y=54.0), Point(x=193.5, y=128.0),  Point(x=354.375, y=42.0), Point(x=353.25, y=40.0), Point(x=351.0, y=40.0), Point(x=349.875, y=38.0), Point(x=347.625, y=38.0), Point(x=346.5, y=36.0)], class_id=0, detection_id='d5c78348-38e1-4281-aa68-9edcbf2cad9e', parent_id=None),
InstanceSegmentationPrediction(x=367.5, y=536.0, width=43.0, height=38.0, confidence=0.8523976802825928, class_name='moeda', class_confidence=None, points=[Point(x=354.375, y=518.0), Point(x=353.25, y=520.0), Point(x=351.0, y=520.0), Point(x=352.125, y=550.0), Point(x=357.75, y=550.0), Point(x=360.0, y=554.0), Point(x=374.625, y=554.0), Point(x=375.75, y=552.0), Point(x=376.875, y=552.0),Point(x=381.375, y=518.0)], class_id=1, detection_id='c93327f3-afce-4038-932b-1fc623fcc949', parent_id=None)])]

However, I need to reorganize my data so that it looks like this:

{'predictions': [{'x': 352.0, 'y': 565.0, 'width': 420.0, 'height': 1060.0, 'confidence': 0.9052466154098511, 'class': 'animal', 'points': [{'x': 244.125, 'y': 36.0}, {'x': 243.0, 'y': 38.0}, {'x': 241.875, 'y': 38.0}, {'x': 239.625, 'y': 42.0}, {'x': 238.5, 'y': 42.0}, {'x': 232.875, 'y': 52.0}, {'x': 232.875, 'y': 54.0}, {'x': 229.5, 'y': 60.0}, {'x': 229.5, 'y': 62.0}, {'x': 209.25, }, {'x': 354.375, 'y': 42.0}, {'x': 353.25, 'y': 40.0}, {'x': 351.0, 'y': 40.0}, {'x': 349.875, 'y': 38.0}, {'x': 348.75, 'y': 38.0}, {'x': 347.625, 'y': 36.0}], 
'class_id': 0, 'detection_id': '71c453d6-3654-4a53-a56c-4db6a012cfe8', 
'image_path': '/content/17.jpeg', 'prediction_type': 'InstanceSegmentationModel'}, {'x': 368.0, 'y': 536.0, 'width': 44.0, 'height': 38.0, 'confidence': 0.8534654378890991, 'class': 'moeda', 'points': [{'x': 354.375, 'y': 518.0}, {'x': 353.25, 'y': 520.0}, {'x': 351.0, 'y': 520.0}, {'x': 348.75, 'y': 524.0}, {'x': 347.625, 'y': 524.0}, {'x': 346.5, 'y': 526.0}, {'x': 346.5, 'y': 542.0}, {'x': 348.75, 'y': 542.0}, {'x': 351.0, 'y': 546.0}, {'x': 351.0,  {'x': 387.0, 'y': 530.0}, {'x': 385.875, 'y': 528.0}, {'x': 385.875, 'y': 526.0}, {'x': 384.75, 'y': 524.0}, {'x': 384.75, 'y': 522.0}, {'x': 383.625, 'y': 522.0}, {'x': 381.375, 'y': 518.0}], 'class_id': 1, 'detection_id': 'a37bfa9e-c169-4d4b-b46e-213ad9508d9b', 'image_path': '/content/17.jpeg', 'prediction_type': 'InstanceSegmentationModel'}], 'image': {'width': '720', 'height': '1280'}}

With this, I can access the points by doing:

    if result['predictions']: 
      points_1 = result['predictions'][1]['points']
      points_o = result['predictions'][0]['points']

I need to access the points of the first data, I tried to change InstanceSegmentationInferenceResponse to the specific dictionary but I couldn't do it.


Solution

  • I'm assuming you're getting this as a response from the inference call on Roboflow and you want to use a more JSON-esque approach to things, I used (assuming you set your response to 'model_response'):

    json_data_string = model_response.model_dump_json()
    results = json.loads(json_data_string)
    

    Some real good documentation for how it's set up is here.