如何在Python中将实例分割预测结果转换为自定义字典格式？

问题描述 #

使用RoboFlow训练的实例分割模型得到的预测结果为InstanceSegmentationInferenceResponse对象列表，格式示例如下：

[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), ...], 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), ...], class_id=1, detection_id='c93327f3-afce-4038-932b-1fc623fcc949', parent_id=None)])]

需要将其转换为指定字典格式，以便通过result['predictions'][0]['points']的方式直接访问关键点。

解决方案 #

直接提取RoboFlow响应对象的公开属性，将对象结构转换为目标字典格式，代码如下：

转换函数 #

def convert_roboflow_response(response_list, image_path="/content/17.jpeg"):
    # 取列表中的第一个推理响应对象
    response = response_list[0]
    converted_predictions = []
    
    for pred in response.predictions:
        # 将Point对象列表转换为字典列表
        points = [{"x": point.x, "y": point.y} for point in pred.points]
        
        # 组装单个预测结果的字典
        pred_dict = {
            "x": pred.x,
            "y": pred.y,
            "width": pred.width,
            "height": pred.height,
            "confidence": pred.confidence,
            "class": pred.class_name,
            "points": points,
            "class_id": pred.class_id,
            "detection_id": pred.detection_id,
            "image_path": image_path,
            "prediction_type": "InstanceSegmentationModel"
        }
        converted_predictions.append(pred_dict)
    
    # 组装最终结果字典
    result = {
        "predictions": converted_predictions,
        "image": {
            "width": str(response.image.width),
            "height": str(response.image.height)
        }
    }
    return result

使用示例 #

# 假设raw_response是你从RoboFlow得到的原始响应列表
raw_response = [InstanceSegmentationInferenceResponse(visualization=None, frame_id=None, time=None, image=InferenceResponseImage(width=720, height=1280), predictions=[...])]

# 执行转换
formatted_result = convert_roboflow_response(raw_response)

# 访问关键点
if formatted_result['predictions']:
    points_o = formatted_result['predictions'][0]['points']  # 第一个实例的关键点
    points_1 = formatted_result['predictions'][1]['points']  # 第二个实例的关键点

说明 #

RoboFlow的推理响应对象（InstanceSegmentationInferenceResponse、InstanceSegmentationPrediction、Point）均提供公开的属性访问，直接通过.属性名即可获取对应值，无需复杂的序列化操作。上述代码仅需遍历提取所需字段，补充指定的额外字段（image_path、prediction_type），即可得到目标格式的字典。

内容的提问来源于stack exchange，提问作者user19270359