From 0343f16c4a3ca2f8e5a45c138685035929580681 Mon Sep 17 00:00:00 2001 From: Greg Heinrich Date: Tue, 23 Aug 2016 14:19:56 +0200 Subject: [PATCH] Allow segmentation color map to be specified explicitly This is useful when working with image labels in RGB format. A separate text file can be provided to specify color/class mappings. Conflicts: digits/extensions/data/imageSegmentation/forms.py --- .../extensions/data/imageSegmentation/data.py | 43 ++++++++++---- .../data/imageSegmentation/forms.py | 10 +++- .../data/imageSegmentation/template.html | 57 +++++++++++++------ 3 files changed, 80 insertions(+), 30 deletions(-) diff --git a/digits/extensions/data/imageSegmentation/data.py b/digits/extensions/data/imageSegmentation/data.py index 437024afe..c9d7df53a 100644 --- a/digits/extensions/data/imageSegmentation/data.py +++ b/digits/extensions/data/imageSegmentation/data.py @@ -7,6 +7,7 @@ import numpy as np import PIL.Image +import PIL.ImagePalette from digits.utils import image, subclass, override, constants from digits.utils.constants import COLOR_PALETTE_ATTRIBUTE @@ -15,7 +16,6 @@ TEMPLATE = "template.html" - @subclass class DataIngestion(DataIngestionInterface): """ @@ -37,12 +37,26 @@ def __init__(self, **kwargs): random.seed(self.userdata['seed']) - # open first image in label folder to retrieve palette - # all label images must use the same palette - this is enforced - # during dataset creation - filename = self.make_image_list(self.label_folder)[0] - image = self.load_label(filename) - self.userdata[COLOR_PALETTE_ATTRIBUTE] = image.getpalette() + if self.userdata['colormap_method'] == "label": + # open first image in label folder to retrieve palette + # all label images must use the same palette - this is enforced + # during dataset creation + filename = self.make_image_list(self.label_folder)[0] + image = self.load_label(filename) + self.userdata[COLOR_PALETTE_ATTRIBUTE] = image.getpalette() + else: + # read colormap from file + with open(self.colormap_text_file) as f: + palette = [] + lines = f.read().splitlines() + for line in lines: + for val in line.split(): + palette.append(int(val)) + # fill rest with zeros + palette = palette + [0] * (256*3 - len(palette)) + self.userdata[COLOR_PALETTE_ATTRIBUTE] = palette + self.palette_img = PIL.Image.new("P", (1,1)) + self.palette_img.putpalette(palette) # get labels if those were provided if self.class_labels_file: @@ -168,10 +182,17 @@ def load_label(self, filename): Load a label image """ image = PIL.Image.open(filename) - if image.mode not in ['P', 'L', '1']: - raise ValueError("Labels are expected to be single-channel (paletted or " - " grayscale) images - %s mode is '%s'" - % (filename, image.mode)) + if self.userdata['colormap_method'] == "label": + if image.mode not in ['P', 'L', '1']: + raise ValueError("Labels are expected to be single-channel (paletted or " + " grayscale) images - %s mode is '%s'" + % (filename, image.mode)) + else: + if image.mode not in ['RGB']: + raise ValueError("Labels are expected to be RGB images - %s mode is '%s'" + % (filename, image.mode)) + image = image.quantize(palette=self.palette_img) + return image def make_image_list(self, folder): diff --git a/digits/extensions/data/imageSegmentation/forms.py b/digits/extensions/data/imageSegmentation/forms.py index beed5ee8a..0a9e26148 100644 --- a/digits/extensions/data/imageSegmentation/forms.py +++ b/digits/extensions/data/imageSegmentation/forms.py @@ -58,7 +58,9 @@ def validate_file_path(form, field): " image folder there must be one corresponding image in the label" " image folder. The label image must have the same filename except" " for the extension, which may differ. Label images are expected" - " to be single-channel images (paletted or grayscale)." + " to be single-channel images (paletted or grayscale), or RGB" + " images, in which case the color/class mappings need to be" + " specified through a separate text file." ) folder_pct_val = utils.forms.IntegerField( @@ -94,7 +96,9 @@ def validate_file_path(form, field): " image folder there must be one corresponding image in the label" " image folder. The label image must have the same filename except" " for the extension, which may differ. Label images are expected" - " to be single-channel images (paletted or grayscale)." + " to be single-channel images (paletted or grayscale), or RGB" + " images, in which case the color/class mappings need to be" + " specified through a separate text file." ) channel_conversion = utils.forms.SelectField( @@ -115,7 +119,7 @@ def validate_file_path(form, field): validate_file_path, ], tooltip="The 'i'th line of the file should give the string label " - "associated with the '(i-1)'th numberic label. (E.g. the " + "associated with the '(i-1)'th numeric label. (E.g. the " "string label for the numeric label 0 is supposed to be " "on line 1.)" ) diff --git a/digits/extensions/data/imageSegmentation/template.html b/digits/extensions/data/imageSegmentation/template.html index be676c7dc..1874a16c3 100644 --- a/digits/extensions/data/imageSegmentation/template.html +++ b/digits/extensions/data/imageSegmentation/template.html @@ -7,13 +7,13 @@
{{ form.feature_folder.label }} {{ form.feature_folder.tooltip }} - {{ form.feature_folder(class='form-control autocomplete_path', placeholder='folder or URL')}} + {{ form.feature_folder(class='form-control autocomplete_path', placeholder='folder')}}
{{ form.label_folder.label }} {{ form.label_folder.tooltip }} - {{ form.label_folder(class='form-control autocomplete_path', placeholder='folder or URL')}} + {{ form.label_folder(class='form-control autocomplete_path', placeholder='folder')}}
@@ -28,13 +28,37 @@ + +
+ {{ form.class_labels_file.label }} + {{ form.class_labels_file.tooltip }} + {{ form.class_labels_file(class='form-control autocomplete_path', placeholder='file')}} +
+ +
+ {{ form.colormap_method.label }} + {{ form.colormap_method.tooltip }} + {{ form.colormap_method(class='form-control') }} +
+ +
+ {{ form.colormap_text_file.label }} + {{ form.colormap_text_file.tooltip }} + {{ form.colormap_text_file(class='form-control autocomplete_path', placeholder='file')}} +
+ +
+ {{ form.channel_conversion.label }} + {{ form.channel_conversion.tooltip }} + {{ form.channel_conversion(class='form-control') }}
-
- {{ form.class_labels_file.label }} - {{ form.class_labels_file.tooltip }} - {{ form.class_labels_file(class='form-control autocomplete_path', placeholder='folder or URL')}} -
- -
- {{ form.channel_conversion.label }} - {{ form.channel_conversion.tooltip }} - {{ form.channel_conversion(class='form-control') }} -
+function hasColorMapMethodChanged() { + if ($("#colormap_method").val() == "textfile") + { + $("#colormap_text_file").prop("disabled", false); + } + else + { + $("#colormap_text_file").prop("disabled", true); + } +} +$("#colormap_method").click(hasColorMapMethodChanged); +hasColorMapMethodChanged(); +