Neodroid  0.2.0
Machine Learning Environment Prototyping Tool
CategorySegmenter.cs
Go to the documentation of this file.
1 using System;
2 using System.Collections.Generic;
5 using UnityEngine;
6 
7 namespace droid.Runtime.Utilities.GameObjects.NeodroidCamera.Segmentation {
12  Tag_,
13  Layer_
14  }
15 
19  [ExecuteInEditMode]
20  public class CategorySegmenter : Segmenter {
23  Renderer[] _all_renders = null;
24 
27  MaterialPropertyBlock _block = null;
28 
29  [SerializeField] Shader segmentation_shader = null;
30  [SerializeField] Camera _camera = null;
31 
32  [SerializeField] protected ColorByCategory[] _colors_by_category = null;
33 
34  [SerializeField] SegmentationMode _segmentation_mode = SegmentationMode.Tag_;
35 
36  [SerializeField] ScriptableObjects.Segmentation _segmentation_preset = null;
37 
40  public bool _Replace_Untagged_Color = true;
41 
44  public Color _Untagged_Color = Color.black;
45 
48  public ColorByCategory[] ColorsByCategory { get { return this._colors_by_category; } }
49 
52  public Dictionary<string, Color> ColorsDictGameObject { get; set; } = new Dictionary<string, Color>();
53 
56  public override Dictionary<string, Color> ColorsDict {
57  get {
58  var colors = new Dictionary<string, Color>();
59  foreach (var key_val in this.ColorsDictGameObject) {
60  colors.Add(key_val.Key, key_val.Value);
61  }
62 
63  return colors;
64  }
65  }
66 
69  void Start() {
70  //this.Setup();
71  }
72 
75  void Awake() { this.Setup(); }
76 
77  void CheckBlock() {
78  if (this._block == null) {
79  this._block = new MaterialPropertyBlock();
80  }
81  }
82 
83  SynthesisUtilities.CapturePass[] _capture_passes;
84 
87  void Setup() {
88  if (this._colors_by_category != null && this._colors_by_category.Length > 0) {
89  foreach (var tag_color in this._colors_by_category) {
90  if (!this.ColorsDictGameObject.ContainsKey(tag_color._Category_Name)) {
91  this.ColorsDictGameObject.Add(tag_color._Category_Name, tag_color._Color);
92  }
93  }
94  }
95 
96  if (this._segmentation_preset) {
97  var segmentation_color_by_tags = this._segmentation_preset._color_by_categories;
98  if (segmentation_color_by_tags != null) {
99  foreach (var tag_color in segmentation_color_by_tags) {
100  if (!this.ColorsDictGameObject.ContainsKey(tag_color._Category_Name)) {
101  this.ColorsDictGameObject.Add(tag_color._Category_Name, tag_color._Color);
102  }
103  }
104  }
105  }
106 
107  this._all_renders = FindObjectsOfType<Renderer>();
108  if (!this._camera) {
109  this._camera = this.GetComponent<Camera>();
110  }
111 
112  if (this.ColorsDictGameObject == null) {
113  this.ColorsDictGameObject = new Dictionary<string, Color>();
114  }
115 
116  switch (this._segmentation_mode) {
117  case SegmentationMode.Tag_:
118  this._capture_passes = new[] {
119  new SynthesisUtilities.CapturePass {
120  _Name = "_tag_id",
121  _ReplacementMode =
122  SynthesisUtilities
123  .ReplacementModes
124  .Tag_id_,
125  _SupportsAntialiasing =
126  false
127  }
128  };
129  break;
130  case SegmentationMode.Layer_:
131  this._capture_passes = new[] {
132  new SynthesisUtilities.CapturePass {
133  _Name = "_layer_id",
134  _ReplacementMode =
135  SynthesisUtilities
136  .ReplacementModes
137  .Layer_id_,
138  _SupportsAntialiasing =
139  false
140  }
141  };
142  break;
143  default: throw new ArgumentOutOfRangeException();
144  }
145 
146  SynthesisUtilities.SetupCapturePassesReplacementShader(this._camera,
147  this.segmentation_shader,
148  ref this._capture_passes);
149 
150  this.CheckBlock();
151  foreach (var a_renderer in this._all_renders) {
152  a_renderer.GetPropertyBlock(this._block);
153  string category_name;
154  var category_int = 0;
155  Color color;
156  string shader_data_name;
157  switch (this._segmentation_mode) {
158  case SegmentationMode.Tag_:
159  category_name = a_renderer.tag;
160  shader_data_name = SynthesisUtilities._Shader_Tag_Color_Name;
161  break;
162  case SegmentationMode.Layer_:
163  category_int = a_renderer.gameObject.layer;
164  category_name = LayerMask.LayerToName(category_int);
165  shader_data_name = SynthesisUtilities._Shader_Layer_Color_Name;
166  break;
167  default: throw new ArgumentOutOfRangeException();
168  }
169 
170  if (!this.ColorsDictGameObject.ContainsKey(category_name)) {
171  switch (this._segmentation_mode) {
172  case SegmentationMode.Tag_:
173  category_int = category_name.GetHashCode();
174  color = ColorEncoding.EncodeTagHashCodeAsColor(category_int);
175  //color = ColorEncoding.EncodeIdAsColor(category_int);
176  break;
177  case SegmentationMode.Layer_:
178  color = ColorEncoding.EncodeLayerAsColor(category_int);
179  break;
180  default:
181  //color = ColorEncoding.EncodeIdAsColor(category_int);
182  throw new ArgumentOutOfRangeException();
183  }
184 
185  this.ColorsDictGameObject.Add(category_name, color);
186  } else {
187  color = this.ColorsDictGameObject[category_name];
188  }
189 
190  this._block.SetColor(shader_data_name, color);
191 
192  a_renderer.SetPropertyBlock(this._block);
193  }
194  }
195  }
196 }