From 0124c5c64fbe9b9fc6a2db380ff979a348c97669 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Mon, 30 Aug 2021 05:19:08 -0500
Subject: [PATCH] Add IndirectArray(seg)

This makes it easy to visualize the results of segmentation.
---
 Project.toml             |  2 ++
 src/ImageSegmentation.jl |  4 +++-
 src/core.jl              | 24 ++++++++++++++++++++++++
 test/core.jl             |  4 ++++
 test/runtests.jl         |  5 +++--
 5 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index d0665ee..752f3d4 100644
--- a/Project.toml
+++ b/Project.toml
@@ -10,6 +10,7 @@ Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
 ImageFiltering = "6a3955dd-da59-5b1f-98d4-e7296123deb5"
 ImageMorphology = "787d08f9-d448-5407-9aad-5290dd7ab264"
+IndirectArrays = "9b13fd28-a010-5f03-acff-a1bbcff69959"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MetaGraphs = "626554b9-1ddb-594c-aa3c-2596fe9399a5"
@@ -27,6 +28,7 @@ Documenter = "0.24, 0.25"
 ImageCore = "0.8.6"
 ImageFiltering = "0.6"
 ImageMorphology = "0.2.6"
+IndirectArrays = "1"
 LightGraphs = "1.1"
 MetaGraphs = "0.6.6"
 RegionTrees = "0.2, 0.3"
diff --git a/src/ImageSegmentation.jl b/src/ImageSegmentation.jl
index 8032278..3108a81 100644
--- a/src/ImageSegmentation.jl
+++ b/src/ImageSegmentation.jl
@@ -4,6 +4,8 @@ import Base: show
 
 using LinearAlgebra, Statistics
 using DataStructures, StaticArrays, ImageCore, ImageFiltering, ImageMorphology, LightGraphs, SimpleWeightedGraphs, RegionTrees, Distances, StaticArrays, Clustering, MetaGraphs
+using ColorVectorSpace: MathTypes
+using IndirectArrays
 import Clustering: kmeans, fuzzy_cmeans
 
 const PairOrTuple{K,V} = Union{Pair{K,V},Tuple{K,V}}
@@ -43,7 +45,7 @@ export
     kmeans,
     fuzzy_cmeans,
     merge_segments,
-    
+
     # types
     SegmentedImage,
     ImageEdge
diff --git a/src/core.jl b/src/core.jl
index 6fa908c..c0ba198 100644
--- a/src/core.jl
+++ b/src/core.jl
@@ -72,6 +72,30 @@ function show(io::IO, seg::SegmentedImage)
     print(io, "Segmented Image with:\n  labels map: ", summary(labels_map(seg)), "\n  number of labels: ", length(segment_labels(seg)))
 end
 
+"""
+    IndirectArray(seg::SegmentedImage, label_values::AbstractDict)
+
+Create an array where index `i...` has value `label_values[labels_map[seg][i...]]`.
+To display each segment using its mean value, use `IndirectArray(seg, segment_mean(seg))`.
+"""
+IndirectArrays.IndirectArray(seg::SegmentedImage, label_values::AbstractDict) = IndirectArray(seg.image_indexmap, label_values)
+
+"""
+    IndirectArray(f, seg::SegmentedImage)
+    IndirectArray(seg::SegmentedImage)
+
+Create a colorized image from `seg`, for visualization of the segments. The colors are chosen by `Colors.distinguishable_colors`.
+
+Optionally pass a function `f` to transform the dictionary before construction.
+The default value is `f=identity`, but `f=OrderedCollections.freeze` is strongly recommended when the number of distinct
+segments is small.
+"""
+function IndirectArrays.IndirectArray(f, seg::SegmentedImage)
+    colors = f(Dict(zip(seg.segment_labels, distinguishable_colors(length(seg.segment_labels)))))
+    return IndirectArray(seg, colors)
+end
+IndirectArrays.IndirectArray(seg::SegmentedImage) = IndirectArray(identity, seg)
+
 """
     G, vert_map = region_adjacency_graph(seg, weight_fn)
 
diff --git a/test/core.jl b/test/core.jl
index 1e8a43d..b5a0ac2 100644
--- a/test/core.jl
+++ b/test/core.jl
@@ -5,6 +5,10 @@
   img[1:2,1:2] .= 2.0
   img[1:2,3:4] .= 3.0
   seg = fast_scanning(img, 0.5)
+  lbl = IndirectArray(seg)
+  @test lbl[1,1] != lbl[1,4] != lbl[4,4]
+  lbl = IndirectArray(freeze, seg)
+  @test lbl[1,1] != lbl[1,4] != lbl[4,4]
 
   @test labels_map(seg) == seg.image_indexmap
   @test segment_labels(seg) == seg.segment_labels
diff --git a/test/runtests.jl b/test/runtests.jl
index a8c257c..12e45aa 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,5 +1,6 @@
-using ImageSegmentation, Images, Test, SimpleWeightedGraphs, LightGraphs, StaticArrays, DataStructures
-using RegionTrees: isleaf, Cell, split! 
+using ImageSegmentation, Images, IndirectArrays, Test, SimpleWeightedGraphs, LightGraphs, StaticArrays, DataStructures
+using DataStructures.OrderedCollections
+using RegionTrees: isleaf, Cell, split!
 using MetaGraphs: MetaGraph, clear_props!, get_prop, has_prop, set_prop!, props, vertices
 
 using Documenter