diff --git a/app.py b/app.py index 09cb227..4b2b057 100644 --- a/app.py +++ b/app.py @@ -301,15 +301,41 @@ async def detect(websocket: WebSocket): "box": [0, 0, 0, 0] } - # Convert detection format to match protocol - flatten detection attributes - detection_dict = {} + # Convert detection format to match backend expectations exactly as in worker.md section 4.2 + detection_dict = { + "carModel": None, + "carBrand": None, + "carYear": None, + "bodyType": None, + "licensePlateText": None, + "licensePlateConfidence": None + } # Handle different detection result formats if isinstance(highest_confidence_detection, dict): - # Copy all fields from the detection result - for key, value in highest_confidence_detection.items(): - if key not in ["box", "id"]: # Skip internal fields - detection_dict[key] = value + # Extract and flatten branch results from parallel classification + branch_results = highest_confidence_detection.get("branch_results", {}) + if branch_results: + logger.debug(f"Processing branch results: {branch_results}") + + # Transform branch results into backend-expected detection attributes + for branch_id, branch_data in branch_results.items(): + if isinstance(branch_data, dict): + # Map common classification fields to backend-expected names + if "brand" in branch_data: + detection_dict["carBrand"] = branch_data["brand"] + if "body_type" in branch_data: + detection_dict["bodyType"] = branch_data["body_type"] + if "class" in branch_data: + class_name = branch_data["class"] + + # Map based on branch/model type + if "brand" in branch_id.lower(): + detection_dict["carBrand"] = class_name + elif "bodytype" in branch_id.lower() or "body" in branch_id.lower(): + detection_dict["bodyType"] = class_name + + logger.info(f"Detection payload: {detection_dict}") detection_data = { "type": "imageDetection",