gimlet-api 0.0.0.dev0__py3-none-any.whl → 0.0.2__py3-none-any.whl

{gimlet_api-0.0.0.dev0.dist-info → gimlet_api-0.0.2.dist-info}/METADATA

@@ -11,6 +11,6 @@ Requires-Dist: protobuf
 Requires-Dist: grpcio
 Requires-Dist: torch
 Requires-Dist: torch_mlir_gml
-Version: 0.0.0.dev0
+Version: 0.0.2
 
 UNKNOWN

{gimlet_api-0.0.0.dev0.dist-info → gimlet_api-0.0.2.dist-info}/RECORD

@@ -35,7 +35,7 @@ gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2.py,sha256=2yD8ZKS5KM
 gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2_grpc.py,sha256=q1PugN3Jm_4v5hVWADJLCIkIEC2_beKEqEH4vb_SpH8,7396
 gml/proto/src/controlplane/model/mpb/v1/mpb_pb2.py,sha256=RVedXkNYu2iF5OHiXoYyRw9AGRCUWG7qNyY-5QY71Go,3762
 gml/proto/src/controlplane/model/mpb/v1/mpb_pb2_grpc.py,sha256=KSdb6V04qUHDsb1R2o3wixwTyZgrhwnPYobjnRgWX4I,4735
-gml/tensor.py,sha256=URYDHDwcbgxdGXhhvKrKelEU_x5UvPdKonWBm8YHEwY,6119
-gimlet_api-0.0.0.dev0.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
-gimlet_api-0.0.0.dev0.dist-info/METADATA,sha256=n4OP-wBQzUZiPqXhY9UUW7t1WwZZaRTlClgcrweImvE,434
-gimlet_api-0.0.0.dev0.dist-info/RECORD,,
+gml/tensor.py,sha256=veEDZGWRCJGa16gAabuCZwSS3jLXDXBk4xTH-v5C-Dw,7170
+gimlet_api-0.0.2.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
+gimlet_api-0.0.2.dist-info/METADATA,sha256=-XADMEHG4d3Gk1up3KH_9NhGK_2qtZG3dbC6NEONG0Q,429
+gimlet_api-0.0.2.dist-info/RECORD,,

gml/tensor.py

@@ -14,7 +14,7 @@
 # SPDX-License-Identifier: Proprietary
 
 import abc
-from typing import List, Optional, Tuple
+from typing import List, Literal, Optional, Tuple
 
 import gml.proto.src.api.corepb.v1.model_exec_pb2 as modelexecpb
 
@@ -156,6 +156,33 @@ class DetectionOutputDimension(DimensionSemantics):
         )
 
 
+def _segmentation_mask_kind_to_proto(kind: str):
+    match kind.lower():
+        case "bool_masks":
+            return (
+                modelexecpb.DimensionSemantics.SegmentationMaskParams.SEGMENTATION_MASK_KIND_BOOL
+            )
+        case "int_label_masks":
+            return (
+                modelexecpb.DimensionSemantics.SegmentationMaskParams.SEGMENTATION_MASK_KIND_CLASS_LABEL
+            )
+        case _:
+            raise ValueError("Invalid segmentation mask kind: {}".format(kind))
+
+
+class SegmentationMaskChannel(DimensionSemantics):
+    def __init__(self, kind: Literal["bool_masks", "int_label_masks"]):
+        self.kind = _segmentation_mask_kind_to_proto(kind)
+
+    def to_proto(self) -> modelexecpb.DimensionSemantics:
+        return modelexecpb.DimensionSemantics(
+            kind=modelexecpb.DimensionSemantics.DIMENSION_SEMANTICS_KIND_SEGMENTATION_MASK_CHANNEL,
+            segmentation_mask_params=modelexecpb.DimensionSemantics.SegmentationMaskParams(
+                kind=self.kind,
+            ),
+        )
+
+
 class TensorSemantics:
     def __init__(self, dimensions: List[DimensionSemantics]):
         self.dimensions = dimensions