Skip to content

Commit

Permalink
Allow segmentation color map to be specified explicitly
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@gheinrich
gheinrich committed Aug 30, 2016
1 parent c88d568 commit 0343f16
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 30 deletions.
43 changes: 32 additions & 11 deletions digits/extensions/data/imageSegmentation/data.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -15,7 +16,6 @@

TEMPLATE = "template.html"


@subclass
class DataIngestion(DataIngestionInterface):
"""
Expand All @@ -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:
Expand Down Expand Up @@ -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):
Expand Down
10 changes: 7 additions & 3 deletions digits/extensions/data/imageSegmentation/forms.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand All @@ -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.)"
)
57 changes: 41 additions & 16 deletions digits/extensions/data/imageSegmentation/template.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,13 @@
<div class="form-group{{mark_errors([form.feature_folder])}}">
{{ 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')}}
</div>

<div class="form-group{{mark_errors([form.label_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')}}
</div>

<div class="form-group{{mark_errors([form.folder_pct_val])}}">
Expand All @@ -28,13 +28,37 @@
<div style="display:none;" class="form-group{{mark_errors([form.validation_feature_folder])}} cl-separate-val-folder">
{{ form.validation_feature_folder.label }}
{{ form.validation_feature_folder.tooltip }}
{{ form.validation_feature_folder(class='form-control autocomplete_path') }}
{{ form.validation_feature_folder(class='form-control autocomplete_path', placeholder='folder') }}
</div>

<div style="display:none;" class="form-group{{mark_errors([form.validation_label_folder])}} cl-separate-val-folder">
{{ form.validation_label_folder.label }}
{{ form.validation_label_folder.tooltip }}
{{ form.validation_label_folder(class='form-control autocomplete_path') }}
{{ form.validation_label_folder(class='form-control autocomplete_path', placeholder='folder') }}
</div>

<div class="form-group{{mark_errors([form.class_labels_file])}}">
{{ form.class_labels_file.label }}
{{ form.class_labels_file.tooltip }}
{{ form.class_labels_file(class='form-control autocomplete_path', placeholder='file')}}
</div>

<div class="form-group{{mark_errors([form.colormap_method])}}">
{{ form.colormap_method.label }}
{{ form.colormap_method.tooltip }}
{{ form.colormap_method(class='form-control') }}
</div>

<div class="form-group{{mark_errors([form.colormap_text_file])}} cl-colormap-text-file">
{{ form.colormap_text_file.label }}
{{ form.colormap_text_file.tooltip }}
{{ form.colormap_text_file(class='form-control autocomplete_path', placeholder='file')}}
</div>

<div class="form-group{{mark_errors([form.channel_conversion])}}">
{{ form.channel_conversion.label }}
{{ form.channel_conversion.tooltip }}
{{ form.channel_conversion(class='form-control') }}
</div>

<script>
Expand All @@ -56,16 +80,17 @@
}
$("#has_val_folder").click(hasValFolderChanged);
hasValFolderChanged();
</script>

<div class="form-group{{mark_errors([form.class_labels_file])}}">
{{ form.class_labels_file.label }}
{{ form.class_labels_file.tooltip }}
{{ form.class_labels_file(class='form-control autocomplete_path', placeholder='folder or URL')}}
</div>

<div class="form-group{{mark_errors([form.channel_conversion])}}">
{{ form.channel_conversion.label }}
{{ form.channel_conversion.tooltip }}
{{ form.channel_conversion(class='form-control') }}
</div>
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();
</script>

0 comments on commit 0343f16

Please sign in to comment.