File size: 4,608 Bytes
960a64d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Response dataclasses for Pete."""

import dataclasses
import enum
from typing import Any, List, Mapping, Optional, Sequence

from ez_wsi_dicomweb import patch_embedding_endpoints

import pete_errors
from data_models import patch_coordinate

_MAX_ERROR_DESCRIPTION_LENGTH = 1024


class ErrorCode(enum.Enum):
  """The error codes for PeteErrorResponse mapped from PeteErrors."""

  TOO_MANY_PATCHES_ERROR = 'TOO_MANY_PATCHES_ERROR'
  INVALID_CREDENTIALS_ERROR = (
      patch_embedding_endpoints.EndpointJsonKeys.INVALID_CREDENTIALS
  )
  PATCH_DIMENSIONS_DO_NOT_MATCH_ENDPOINT_INPUT_DIMENSIONS_ERROR = (
      'PATCH_DIMENSIONS_DO_NOT_MATCH_ENDPOINT_INPUT_DIMENSIONS_ERROR'
  )
  INSTANCES_NOT_CONCATENATED_ERROR = 'INSTANCES_NOT_CONCATENATED_ERROR'
  INVALID_REQUEST_FIELD_ERROR = 'INVALID_REQUEST_FIELD_ERROR'
  INVALID_RESPONSE_ERROR = 'INVALID_RESPONSE_ERROR'
  LEVEL_NOT_FOUND_ERROR = 'LEVEL_NOT_FOUND_ERROR'
  EZ_WSI_STATE_ERROR = 'EZ_WSI_STATE_ERROR'
  IMAGE_ERROR = 'IMAGE_ERROR'
  HTTP_ERROR = 'HTTP_ERROR'
  INVALID_ICC_PROFILE_TRANSFORM_ERROR = 'INVALID_ICC_PROFILE_TRANSFORM_ERROR'
  IMAGE_DIMENSION_ERROR = 'IMAGE_DIMENSION_ERROR'
  DICOM_TILED_FULL_ERROR = 'DICOM_TILED_FULL_ERROR'
  DICOM_ERROR = 'DICOM_ERROR'
  DICOM_IMAGE_DOWNSAMPLING_TOO_LARGE_ERROR = (
      'DICOM_IMAGE_DOWNSAMPLING_TOO_LARGE_ERROR'
  )
  PATCH_OUTSIDE_OF_IMAGE_DIMENSIONS_ERROR = (
      'PATCH_OUTSIDE_OF_IMAGE_DIMENSIONS_ERROR'
  )
  DICOM_PATH_ERROR = 'DICOM_PATH_ERROR'
  GCS_IMAGE_PATH_FORMAT_ERROR = 'GCS_IMAGE_PATH_FORMAT_ERROR'
  UNAPPROVED_DICOM_STORE_ERROR = 'UNAPPROVED_DICOM_STORE_ERROR'
  UNAPPROVED_GCS_BUCKET_ERROR = 'UNAPPROVED_GCS_BUCKET_ERROR'


@dataclasses.dataclass(frozen=True)
class PeteErrorResponse:
  """The response when Pete is unable to successfully complete a request."""

  error_code: ErrorCode


@dataclasses.dataclass(frozen=True)
class PatchEmbeddingV1:
  """A List of embeddings, instance uids, and patch coordinate."""

  embeddings: List[float]
  patch_coordinate: patch_coordinate.PatchCoordinate


@dataclasses.dataclass(frozen=True)
class PatchEmbeddingV2:
  """A List of embeddings, instance uids, and patch coordinate."""

  embedding_vector: List[float]
  patch_coordinate: patch_coordinate.PatchCoordinate


@dataclasses.dataclass(frozen=True)
class EmbeddingResultV1:
  """The response when Pete is able to successfully complete a request."""

  dicom_study_uid: str
  dicom_series_uid: str
  instance_uids: List[str]
  patch_embeddings: List[PatchEmbeddingV1]


@dataclasses.dataclass(frozen=True)
class EmbeddingResponseV1:
  """An instance in a Embedding Response as described in the schema file."""

  model_version: str
  error_response: Optional[PeteErrorResponse]
  embedding_result: List[EmbeddingResultV1]

  def __post_init__(self):
    if self.error_response is None and self.embedding_result is None:
      raise pete_errors.InvalidResponseError(
          'At least one of error_response or embedding_result must be set.'
      )


def embedding_instance_response_v2(
    results: Sequence[PatchEmbeddingV2],
) -> Mapping[str, Any]:
  """Returns a JSON-serializable embedding instance responses."""
  return {
      patch_embedding_endpoints.EndpointJsonKeys.RESULT: {
          patch_embedding_endpoints.EndpointJsonKeys.PATCH_EMBEDDINGS: [
              dataclasses.asdict(patch_embedding) for patch_embedding in results
          ]
      },
  }


def instance_error_response_v2(
    error_code: ErrorCode, description: str = ''
) -> Mapping[str, Any]:
  error = {
      patch_embedding_endpoints.EndpointJsonKeys.ERROR_CODE: error_code.value
  }
  if description:
    error[patch_embedding_endpoints.EndpointJsonKeys.ERROR_CODE_DESCRIPTION] = (
        description[:_MAX_ERROR_DESCRIPTION_LENGTH]
    )
  return {
      patch_embedding_endpoints.EndpointJsonKeys.ERROR: error,
  }


def prediction_error_response_v2(error_code: ErrorCode) -> Mapping[str, Any]:
  return {
      patch_embedding_endpoints.EndpointJsonKeys.VERTEXAI_ERROR: (
          error_code.value
      )
  }