Indexed color shader

batch/serializers.py

@@ -3,7 +3,7 @@
 from rest_framework import serializers
 
 from . import models
-import cloud_ilastik.datasets.models as datasets_models
+from cloud_ilastik import datasets
 from cloud_ilastik.datasets import neuroglancer as ng
 
 
@@ -32,7 +32,7 @@ class Meta:
 
 class ResultSerializer(serializers.ModelSerializer):
     class Meta:
-        model = datasets_models.Dataset
+        model = datasets.models.Dataset
         fields = ["neuroglancer_url"]
 
 
@@ -61,7 +61,7 @@ class Meta:
 class BatchJob(serializers.Serializer):
     project = serializers.PrimaryKeyRelatedField(queryset=models.Project.objects.all(), allow_null=False)
     datasets = serializers.PrimaryKeyRelatedField(
-        many=True, queryset=datasets_models.Dataset.objects.all(), allow_null=False
+        many=True, queryset=datasets.models.Dataset.objects.all(), allow_null=False
     )
 
     class Meta:
@@ -72,8 +72,10 @@ class JobUpdate(serializers.Serializer):
     status = serializers.ChoiceField(choices=[models.JobStatus.done.value, models.JobStatus.failed.value])
     result_url = serializers.URLField()
     name = serializers.CharField()
-    dtype = serializers.ChoiceField(choices=datasets_models.DType.values())
-    channel_type = serializers.ChoiceField(choices=datasets_models.ChannelType.values(), default=datasets_models.ChannelType.Intensity.value, required=False)
+    dtype = serializers.ChoiceField(choices=datasets.models.DType.values())
+    channel_type = serializers.ChoiceField(
+        choices=datasets.types.ChannelType.values(), default=datasets.types.ChannelType.Intensity.value, required=False
+    )
     size_t = serializers.IntegerField(default=1)
     size_z = serializers.IntegerField(default=1)
     size_y = serializers.IntegerField()

cloud_ilastik/datasets/apps.py

@@ -2,4 +2,4 @@
 
 
 class DatasetsConfig(AppConfig):
-    name = 'cloud_ilastik.datasets'
+    name = "cloud_ilastik.datasets"

cloud_ilastik/datasets/models.py

@@ -7,6 +7,7 @@
 import files.models as files_models
 
 from . import neuroglancer as ng
+from . import types
 
 TAR_URL_RE = re.compile("/data$")
 
@@ -34,25 +35,13 @@ def values(cls):
         return tuple(item.value for item in cls)
 
 
-@enum.unique
-class ChannelType(enum.Enum):
-    Intensity = "intensity"
-    IndexedColor = "indexed"
-
-    @classmethod
-    def choices(cls):
-        return tuple((item.name, item.value) for item in cls)
-
-    @classmethod
-    def values(cls):
-        return tuple(item.value for item in cls)
-
-
 class Dataset(models.Model):
     name = models.CharField(max_length=255)
     url = models.URLField()
     dtype = models.CharField(max_length=15, choices=DType.choices())
-    channel_type = models.CharField(max_length=15, choices=ChannelType.choices(), default=ChannelType.Intensity.value)
+    channel_type = models.CharField(
+        max_length=15, choices=types.ChannelType.choices(), default=types.ChannelType.Intensity.value
+    )
     size_t = models.PositiveIntegerField(default=1)
     size_z = models.PositiveIntegerField(default=1)
     size_y = models.PositiveIntegerField()
@@ -70,7 +59,6 @@ def save(self, *args, **kwargs):
             self.owner = self.job.owner
         super().save(*args, **kwargs)
 
-
     @property
     def sizes(self):
         return {"t": self.size_t, "z": self.size_z, "y": self.size_y, "x": self.size_x, "c": self.size_c}
@@ -86,7 +74,7 @@ def as_viewer_layer(self):
             mode = ng.ColorMode.RGB
         else:
             mode = ng.ColorMode.ILASTIK
-        return ng.Layer(self.url, self.size_c, color_mode=mode, role="data")
+        return ng.Layer(self.url, self.size_c, color_mode=mode, channel_type=self.channel_type, role="data")
 
     @property
     def neuroglancer_url(self):

cloud_ilastik/datasets/neuroglancer.py

@@ -6,36 +6,34 @@
 
 from django.conf import settings
 
-__all__ = ["viewer_url"]
-
+from . import types
 
-@enum.unique
-class ColorMode(enum.Enum):
-    ILASTIK = "ilastik"
-    RGB = "rgb"
-    GRAYSCALE = "grayscale"
+__all__ = ["viewer_url"]
 
 
 class Layer:
     url: str
     num_channels: int
     role: str
     selected: bool
-    color_mode: ColorMode
+    color_table: types.ColorTable
+    channel_type: types.ChannelType
 
     def __init__(
         self,
         url: str,
         num_channels: int,
         role: str = "data",
         selected: bool = False,
-        color_mode: ColorMode = ColorMode.RGB,
+        color_table: types.ColorTable = types.ColorTable.RGB,
+        channel_type: types.ChannelType = types.ChannelType.Intensity,
     ):
         self.url = url
         self.num_channels = num_channels
         self.role = role
         self.selected = selected
-        self.color_mode = color_mode
+        self.color_table = color_table
+        self.channel_type = channel_type
 
 
 class _Color:
@@ -71,24 +69,24 @@ def as_normalized_vec3(self) -> str:
         return f"vec3({self.r}, {self.g}, {self.b})"
 
     @classmethod
-    def get_colors(cls, num_colors: int, mode: ColorMode) -> List["Color"]:
+    def get_colors(cls, num_colors: int, table: types.ColorTable) -> List["Color"]:
         color_table = {
-            ColorMode.GRAYSCALE: cls.COLORS_GRAYSCALE,
-            ColorMode.RGB: cls.COLORS_RGB,
-            ColorMode.ILASTIK: cls.COLORS_ILASTIK,
-        }[mode]
+            types.ColorTable.GRAYSCALE: cls.COLORS_GRAYSCALE,
+            types.ColorTable.RGB: cls.COLORS_RGB,
+            types.ColorTable.ILASTIK: cls.COLORS_ILASTIK,
+        }[table]
 
         return [cls(*rgb) for rgb in color_table[:num_colors]]
 
     def __repr__(self):
         return f"<Color ({self.r},{self.g},{self.b})>"
 
 
-def _create_fragment_shader(channel_colors: List[_Color]):
+def _create_intensity_fragment_shader(colors):
     color_lines: List[str] = []
     colors_to_mix: List[str] = []
 
-    for idx, color in enumerate(channel_colors):
+    for idx, color in enumerate(colors):
         color_line = f"vec3 color{idx} = ({color.as_normalized_vec3()} / 255.0) * toNormalized(getDataValue({idx}));"
         color_lines.append(color_line)
         colors_to_mix.append(f"color{idx}")
@@ -104,21 +102,46 @@ def _create_fragment_shader(channel_colors: List[_Color]):
     return "\n".join(shader_lines)
 
 
+def _create_indexed_color_fragment_shader(colors):
+    color_lines: List[str] = ["vec4(0.0, 0.0, 0.0, 0.0)"]
+
+    for color in colors:
+        color_lines.append(f"vec4({color.r / 255.0}, {color.g / 255.0}, {color.b / 255.0}, 1.0)")
+
+    return f"""vec4 COLOR_MASKS[{len(color_lines)}] = vec4[](
+    {",".join(color_lines)}
+);
+void main() {{
+  uint val = toRaw(getDataValue());
+  emitRGB(COLOR_MASKS[val]);
+}}"""
+
+
+def _create_fragment_shader(colors: List[_Color], channel_type: types.ChannelType):
+    if channel_type == types.ChannelType.Intensity:
+        return _create_intensity_fragment_shader(colors)
+    elif channel_type == types.ChannelType.IndexedColor:
+        return _create_indexed_color_fragment_shader(colors)
+    else:
+        raise Exception(f"Unknown channel type {channel_type}")
+
+
 def viewer_url(layers: List[Layer], show_control_panel=False) -> str:
     ng_url = "https://web.ilastik.org/viewer/#!"
     ng_layers = []
     selected_layer = None
 
     for layer in layers:
         data_url = layer.url.replace(settings.SWIFT_PREFIX, "https://web.ilastik.org/data/")
+        colors = _Color.get_colors(layer.color_table)
         ng_layers.append(
             {
                 "type": "image",
                 "source": {"url": f"n5://{data_url}"},
                 "tab": "source",
                 "blend": "default",
                 "name": layer.role,
-                "shader": _create_fragment_shader(_Color.get_colors(layer.num_channels, layer.color_mode)),
+                "shader": _create_fragment_shader(colors, layer.channel_type),
             }
         )
 

cloud_ilastik/datasets/types.py

@@ -0,0 +1,22 @@
+import enum
+
+
+@enum.unique
+class ChannelType(enum.Enum):
+    Intensity = "intensity"
+    IndexedColor = "indexed"
+
+    @classmethod
+    def choices(cls):
+        return tuple((item.name, item.value) for item in cls)
+
+    @classmethod
+    def values(cls):
+        return tuple(item.value for item in cls)
+
+
+@enum.unique
+class ColorTable(enum.Enum):
+    ILASTIK = "ilastik"
+    RGB = "rgb"
+    GRAYSCALE = "grayscale"

cloud_ilastik/datasets/urls.py

@@ -2,8 +2,8 @@
 
 from . import views
 
-app_name = 'datasets'
+app_name = "datasets"
 urlpatterns = [
-    path('', views.ListView.as_view(), name='list'),
-    path('<int:pk>/', views.DetailView.as_view(), name='detail'),
+    path("", views.ListView.as_view(), name="list"),
+    path("<int:pk>/", views.DetailView.as_view(), name="detail"),
 ]

tests/test_cloud_ilastik/test_datasets/test_neuroglancer.py

@@ -0,0 +1,14 @@
+from cloud_ilastik.datasets import neuroglancer as ng, types
+
+
+def test_indexed_color_shader():
+    colors = ng._Color.get_colors(None, types.ColorTable.RGB)
+    shader = ng._create_indexed_color_fragment_shader(colors)
+    expected_shader = f"""vec4 COLOR_MASKS[4] = vec4[](
+    vec4(0.0, 0.0, 0.0, 0.0),vec4(1.0, 0.0, 0.0, 1.0),vec4(0.0, 1.0, 0.0, 1.0),vec4(0.0, 0.0, 1.0, 1.0)
+);
+void main() {{
+  uint val = toRaw(getDataValue());
+  emitRGB(COLOR_MASKS[val]);
+}}"""
+    assert shader == expected_shader