nick_93 commited on
Commit
87b4a1a
1 Parent(s): bf2495a
Files changed (4) hide show
  1. app.py +11 -12
  2. colors.py +344 -0
  3. palette.py +37 -0
  4. utils.py +83 -0
app.py CHANGED
@@ -8,13 +8,13 @@ import torch
8
  from diffusers.pipelines.controlnet import StableDiffusionControlNetInpaintPipeline
9
  from diffusers import ControlNetModel, UniPCMultistepScheduler, AutoPipelineForText2Image
10
  from transformers import AutoImageProcessor, UperNetForSemanticSegmentation, AutoModelForDepthEstimation
11
- from models.colors import ade_palette
12
- from models.utils import map_colors_rgb
13
  from diffusers import StableDiffusionXLPipeline
14
  import gradio as gr
15
 
16
- device = "cuda"
17
- dtype = torch.float16
18
 
19
 
20
  css = """
@@ -110,9 +110,9 @@ def segment_image(
110
 
111
  def get_depth_pipeline():
112
  feature_extractor = AutoImageProcessor.from_pretrained("models/models--LiheYoung--depth-anything-large-hf",
113
- torch_dtype=torch.float16)
114
  depth_estimator = AutoModelForDepthEstimation.from_pretrained("models/models--LiheYoung--depth-anything-large-hf",
115
- torch_dtype=torch.float16)
116
  return feature_extractor, depth_estimator
117
 
118
 
@@ -174,16 +174,16 @@ class ControlNetDepthDesignModelMulti:
174
 
175
  #os.environ['HF_HUB_OFFLINE'] = "True"
176
  controlnet_depth= ControlNetModel.from_pretrained(
177
- "models/controlnet_depth", torch_dtype=torch.float16, use_safetensors=True)
178
  controlnet_seg = ControlNetModel.from_pretrained(
179
- "models/own_controlnet", torch_dtype=torch.float16, use_safetensors=True)
180
 
181
  self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
182
  "SG161222/Realistic_Vision_V5.1_noVAE",
183
  #"models/runwayml--stable-diffusion-inpainting",
184
  controlnet=[controlnet_depth, controlnet_seg],
185
  safety_checker=None,
186
- torch_dtype=torch.float16
187
  )
188
 
189
  self.pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models",
@@ -278,7 +278,7 @@ class ControlNetDepthDesignModelMulti:
278
  return design_image
279
 
280
 
281
- def create_refseg_demo(model, device):
282
  gr.Markdown("### Stable Design demo")
283
  with gr.Row():
284
  with gr.Column():
@@ -334,7 +334,6 @@ def create_refseg_demo(model, device):
334
 
335
 
336
  def main():
337
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
338
  model = ControlNetDepthDesignModelMulti()
339
  print('Models uploaded successfully')
340
 
@@ -347,7 +346,7 @@ def main():
347
  gr.Markdown(title)
348
  gr.Markdown(description)
349
 
350
- create_refseg_demo(model, device)
351
  gr.HTML('''<br><br><br><center>You can duplicate this Space to skip the queue:<a href="https://huggingface.co/spaces/MykolaL/StableDesign?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a><br>
352
  <p><img src="https://visitor-badge.glitch.me/badge?page_id=MykolaL/StableDesign" alt="visitors"></p></center>''')
353
 
 
8
  from diffusers.pipelines.controlnet import StableDiffusionControlNetInpaintPipeline
9
  from diffusers import ControlNetModel, UniPCMultistepScheduler, AutoPipelineForText2Image
10
  from transformers import AutoImageProcessor, UperNetForSemanticSegmentation, AutoModelForDepthEstimation
11
+ from colors import ade_palette
12
+ from utils import map_colors_rgb
13
  from diffusers import StableDiffusionXLPipeline
14
  import gradio as gr
15
 
16
+ device = "cpu"#"cuda"
17
+ dtype = torch.float32#torch.float16
18
 
19
 
20
  css = """
 
110
 
111
  def get_depth_pipeline():
112
  feature_extractor = AutoImageProcessor.from_pretrained("models/models--LiheYoung--depth-anything-large-hf",
113
+ torch_dtype=dtype)
114
  depth_estimator = AutoModelForDepthEstimation.from_pretrained("models/models--LiheYoung--depth-anything-large-hf",
115
+ torch_dtype=dtype)
116
  return feature_extractor, depth_estimator
117
 
118
 
 
174
 
175
  #os.environ['HF_HUB_OFFLINE'] = "True"
176
  controlnet_depth= ControlNetModel.from_pretrained(
177
+ "models/controlnet_depth", torch_dtype=dtype, use_safetensors=True)
178
  controlnet_seg = ControlNetModel.from_pretrained(
179
+ "models/own_controlnet", torch_dtype=dtype, use_safetensors=True)
180
 
181
  self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
182
  "SG161222/Realistic_Vision_V5.1_noVAE",
183
  #"models/runwayml--stable-diffusion-inpainting",
184
  controlnet=[controlnet_depth, controlnet_seg],
185
  safety_checker=None,
186
+ torch_dtype=dtype
187
  )
188
 
189
  self.pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models",
 
278
  return design_image
279
 
280
 
281
+ def create_refseg_demo(model):
282
  gr.Markdown("### Stable Design demo")
283
  with gr.Row():
284
  with gr.Column():
 
334
 
335
 
336
  def main():
 
337
  model = ControlNetDepthDesignModelMulti()
338
  print('Models uploaded successfully')
339
 
 
346
  gr.Markdown(title)
347
  gr.Markdown(description)
348
 
349
+ create_refseg_demo(model)
350
  gr.HTML('''<br><br><br><center>You can duplicate this Space to skip the queue:<a href="https://huggingface.co/spaces/MykolaL/StableDesign?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a><br>
351
  <p><img src="https://visitor-badge.glitch.me/badge?page_id=MykolaL/StableDesign" alt="visitors"></p></center>''')
352
 
colors.py ADDED
@@ -0,0 +1,344 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Color mappings"""
2
+ from typing import List
3
+
4
+ TRIVIA = {
5
+ "#B47878": "building;edifice",
6
+ "#06E6E6": "sky",
7
+ "#04C803": "tree",
8
+ "#8C8C8C": "road;route",
9
+ "#04FA07": "grass",
10
+ "#96053D": "person;individual;someone;somebody;mortal;soul",
11
+ "#CCFF04": "plant;flora;plant;life",
12
+ "#787846": "earth;ground",
13
+ "#FF09E0": "house",
14
+ "#0066C8": "car;auto;automobile;machine;motorcar",
15
+ "#3DE6FA": "water",
16
+ "#FF3D06": "railing;rail",
17
+ "#FF5C00": "arcade;machine",
18
+ "#FFE000": "stairs;steps",
19
+ "#00F5FF": "fan",
20
+ "#FF008F": "step;stair",
21
+ "#1F00FF": "stairway;staircase",
22
+ "#FFD600": "radiator",
23
+ }
24
+
25
+ OBJECTS = {
26
+ "#CC05FF": "bed",
27
+ "#FF0633": "painting;picture",
28
+ "#DCDCDC": "mirror",
29
+ "#00FF14": "box",
30
+ "#FF0000": "flower",
31
+ "#FFA300": "book",
32
+ "#00FFC2": "television;television;receiver;television;set;tv;tv;set;idiot;box;boob;tube;telly;goggle;box",
33
+ "#F500FF": "pot;flowerpot",
34
+ "#00FFCC": "vase",
35
+ "#29FF00": "tray",
36
+ "#8FFF00": "poster;posting;placard;notice;bill;card",
37
+ "#5CFF00": "basket;handbasket",
38
+ "#00ADFF": "screen;door;screen",
39
+ }
40
+
41
+
42
+ SITTING = {
43
+ "#0B66FF": "sofa;couch;lounge",
44
+ "#CC4603": "chair",
45
+ "#07FFE0": "seat",
46
+ "#08FFD6": "armchair",
47
+ "#FFC207": "cushion",
48
+ "#00EBFF": "pillow",
49
+ "#00D6FF": "stool",
50
+ "#1400FF": "blanket;cover",
51
+ "#0A00FF": "swivel;chair",
52
+ "#FF9900": "ottoman;pouf;pouffe;puff;hassock",
53
+ }
54
+
55
+ LIGHTING = {
56
+ "#E0FF08": "lamp",
57
+ "#FFAD00": "light;light;source",
58
+ "#001FFF": "chandelier;pendant;pendent",
59
+ }
60
+
61
+ TABLES = {
62
+ "#FF0652": "table",
63
+ "#0AFF47": "desk",
64
+ }
65
+
66
+ CLOSETS = {
67
+ "#E005FF": "cabinet",
68
+ "#FF0747": "shelf",
69
+ "#07FFFF": "wardrobe;closet;press",
70
+ "#0633FF": "chest;of;drawers;chest;bureau;dresser",
71
+ "#0000FF": "case;display;case;showcase;vitrine",
72
+ }
73
+
74
+
75
+ BATHROOM = {
76
+ "#6608FF": "bathtub;bathing;tub;bath;tub",
77
+ "#00FF85": "toilet;can;commode;crapper;pot;potty;stool;throne",
78
+ "#0085FF": "shower",
79
+ "#FF0066": "towel",
80
+ }
81
+
82
+ WINDOWS = {
83
+ "#FF3307": "curtain;drape;drapery;mantle;pall",
84
+ "#E6E6E6": "windowpane;window",
85
+ "#00FF3D": "awning;sunshade;sunblind",
86
+ "#003DFF": "blind;screen",
87
+ }
88
+
89
+ FLOOR = {
90
+ "#FF095C": "rug;carpet;carpeting",
91
+ "#503232": "floor;flooring",
92
+ }
93
+
94
+ INTERIOR = {
95
+ "#787878": "wall",
96
+ "#787850": "ceiling",
97
+ "#08FF33": "door;double;door",
98
+ }
99
+
100
+ KITCHEN = {
101
+ "#00FF29": "kitchen;island",
102
+ "#14FF00": "refrigerator;icebox",
103
+ "#00A3FF": "sink",
104
+ "#EB0CFF": "counter",
105
+ "#D6FF00": "dishwasher;dish;washer;dishwashing;machine",
106
+ "#FF00EB": "microwave;microwave;oven",
107
+ "#47FF00": "oven",
108
+ "#66FF00": "clock",
109
+ "#00FFB8": "plate",
110
+ "#19C2C2": "glass;drinking;glass",
111
+ "#00FF99": "bar",
112
+ "#00FF0A": "bottle",
113
+ "#FF7000": "buffet;counter;sideboard",
114
+ "#B800FF": "washer;automatic;washer;washing;machine",
115
+ "#00FF70": "coffee;table;cocktail;table",
116
+ "#008FFF": "countertop",
117
+ "#33FF00": "stove;kitchen;stove;range;kitchen;range;cooking;stove",
118
+ }
119
+
120
+ LIVINGROOM = {
121
+ "#FA0A0F": "fireplace;hearth;open;fireplace",
122
+ "#FF4700": "pool;table;billiard;table;snooker;table",
123
+ }
124
+
125
+ OFFICE = {
126
+ "#00FFAD": "computer;computing;machine;computing;device;data;processor;electronic;computer;information;processing;system",
127
+ "#00FFF5": "bookcase",
128
+ "#0633FF": "chest;of;drawers;chest;bureau;dresser",
129
+ "#005CFF": "monitor;monitoring;device",
130
+ }
131
+
132
+
133
+ COLOR_MAPPING_CATEGORY_ = {
134
+ 'keep background': {'#FFFFFF': 'background'},
135
+ 'trivia': TRIVIA,
136
+ 'objects': OBJECTS,
137
+ 'sitting': SITTING,
138
+ 'lighting': LIGHTING,
139
+ 'tables': TABLES,
140
+ 'closets': CLOSETS,
141
+ 'bathroom': BATHROOM,
142
+ 'windows': WINDOWS,
143
+ 'floor': FLOOR,
144
+ 'interior': INTERIOR,
145
+ 'kitchen': KITCHEN,
146
+ 'livingroom': LIVINGROOM,
147
+ 'office': OFFICE}
148
+
149
+
150
+ COLOR_MAPPING_ = {
151
+ '#FFFFFF': 'background',
152
+ "#787878": "wall",
153
+ "#B47878": "building;edifice",
154
+ "#06E6E6": "sky",
155
+ "#503232": "floor;flooring",
156
+ "#04C803": "tree",
157
+ "#787850": "ceiling",
158
+ "#8C8C8C": "road;route",
159
+ "#CC05FF": "bed",
160
+ "#E6E6E6": "windowpane;window",
161
+ "#04FA07": "grass",
162
+ "#E005FF": "cabinet",
163
+ "#EBFF07": "sidewalk;pavement",
164
+ "#96053D": "person;individual;someone;somebody;mortal;soul",
165
+ "#787846": "earth;ground",
166
+ "#08FF33": "door;double;door",
167
+ "#FF0652": "table",
168
+ "#8FFF8C": "mountain;mount",
169
+ "#CCFF04": "plant;flora;plant;life",
170
+ "#FF3307": "curtain;drape;drapery;mantle;pall",
171
+ "#CC4603": "chair",
172
+ "#0066C8": "car;auto;automobile;machine;motorcar",
173
+ "#3DE6FA": "water",
174
+ "#FF0633": "painting;picture",
175
+ "#0B66FF": "sofa;couch;lounge",
176
+ "#FF0747": "shelf",
177
+ "#FF09E0": "house",
178
+ "#0907E6": "sea",
179
+ "#DCDCDC": "mirror",
180
+ "#FF095C": "rug;carpet;carpeting",
181
+ "#7009FF": "field",
182
+ "#08FFD6": "armchair",
183
+ "#07FFE0": "seat",
184
+ "#FFB806": "fence;fencing",
185
+ "#0AFF47": "desk",
186
+ "#FF290A": "rock;stone",
187
+ "#07FFFF": "wardrobe;closet;press",
188
+ "#E0FF08": "lamp",
189
+ "#6608FF": "bathtub;bathing;tub;bath;tub",
190
+ "#FF3D06": "railing;rail",
191
+ "#FFC207": "cushion",
192
+ "#FF7A08": "base;pedestal;stand",
193
+ "#00FF14": "box",
194
+ "#FF0829": "column;pillar",
195
+ "#FF0599": "signboard;sign",
196
+ "#0633FF": "chest;of;drawers;chest;bureau;dresser",
197
+ "#EB0CFF": "counter",
198
+ "#A09614": "sand",
199
+ "#00A3FF": "sink",
200
+ "#8C8C8C": "skyscraper",
201
+ "#FA0A0F": "fireplace;hearth;open;fireplace",
202
+ "#14FF00": "refrigerator;icebox",
203
+ "#1FFF00": "grandstand;covered;stand",
204
+ "#FF1F00": "path",
205
+ "#FFE000": "stairs;steps",
206
+ "#99FF00": "runway",
207
+ "#0000FF": "case;display;case;showcase;vitrine",
208
+ "#FF4700": "pool;table;billiard;table;snooker;table",
209
+ "#00EBFF": "pillow",
210
+ "#00ADFF": "screen;door;screen",
211
+ "#1F00FF": "stairway;staircase",
212
+ "#0BC8C8": "river",
213
+ "#FF5200": "bridge;span",
214
+ "#00FFF5": "bookcase",
215
+ "#003DFF": "blind;screen",
216
+ "#00FF70": "coffee;table;cocktail;table",
217
+ "#00FF85": "toilet;can;commode;crapper;pot;potty;stool;throne",
218
+ "#FF0000": "flower",
219
+ "#FFA300": "book",
220
+ "#FF6600": "hill",
221
+ "#C2FF00": "bench",
222
+ "#008FFF": "countertop",
223
+ "#33FF00": "stove;kitchen;stove;range;kitchen;range;cooking;stove",
224
+ "#0052FF": "palm;palm;tree",
225
+ "#00FF29": "kitchen;island",
226
+ "#00FFAD": "computer;computing;machine;computing;device;data;processor;electronic;computer;information;processing;system",
227
+ "#0A00FF": "swivel;chair",
228
+ "#ADFF00": "boat",
229
+ "#00FF99": "bar",
230
+ "#FF5C00": "arcade;machine",
231
+ "#FF00FF": "hovel;hut;hutch;shack;shanty",
232
+ "#FF00F5": "bus;autobus;coach;charabanc;double-decker;jitney;motorbus;motorcoach;omnibus;passenger;vehicle",
233
+ "#FF0066": "towel",
234
+ "#FFAD00": "light;light;source",
235
+ "#FF0014": "truck;motortruck",
236
+ "#FFB8B8": "tower",
237
+ "#001FFF": "chandelier;pendant;pendent",
238
+ "#00FF3D": "awning;sunshade;sunblind",
239
+ "#0047FF": "streetlight;street;lamp",
240
+ "#FF00CC": "booth;cubicle;stall;kiosk",
241
+ "#00FFC2": "television;television;receiver;television;set;tv;tv;set;idiot;box;boob;tube;telly;goggle;box",
242
+ "#00FF52": "airplane;aeroplane;plane",
243
+ "#000AFF": "dirt;track",
244
+ "#0070FF": "apparel;wearing;apparel;dress;clothes",
245
+ "#3300FF": "pole",
246
+ "#00C2FF": "land;ground;soil",
247
+ "#007AFF": "bannister;banister;balustrade;balusters;handrail",
248
+ "#00FFA3": "escalator;moving;staircase;moving;stairway",
249
+ "#FF9900": "ottoman;pouf;pouffe;puff;hassock",
250
+ "#00FF0A": "bottle",
251
+ "#FF7000": "buffet;counter;sideboard",
252
+ "#8FFF00": "poster;posting;placard;notice;bill;card",
253
+ "#5200FF": "stage",
254
+ "#A3FF00": "van",
255
+ "#FFEB00": "ship",
256
+ "#08B8AA": "fountain",
257
+ "#8500FF": "conveyer;belt;conveyor;belt;conveyer;conveyor;transporter",
258
+ "#00FF5C": "canopy",
259
+ "#B800FF": "washer;automatic;washer;washing;machine",
260
+ "#FF001F": "plaything;toy",
261
+ "#00B8FF": "swimming;pool;swimming;bath;natatorium",
262
+ "#00D6FF": "stool",
263
+ "#FF0070": "barrel;cask",
264
+ "#5CFF00": "basket;handbasket",
265
+ "#00E0FF": "waterfall;falls",
266
+ "#70E0FF": "tent;collapsible;shelter",
267
+ "#46B8A0": "bag",
268
+ "#A300FF": "minibike;motorbike",
269
+ "#9900FF": "cradle",
270
+ "#47FF00": "oven",
271
+ "#FF00A3": "ball",
272
+ "#FFCC00": "food;solid;food",
273
+ "#FF008F": "step;stair",
274
+ "#00FFEB": "tank;storage;tank",
275
+ "#85FF00": "trade;name;brand;name;brand;marque",
276
+ "#FF00EB": "microwave;microwave;oven",
277
+ "#F500FF": "pot;flowerpot",
278
+ "#FF007A": "animal;animate;being;beast;brute;creature;fauna",
279
+ "#FFF500": "bicycle;bike;wheel;cycle",
280
+ "#0ABED4": "lake",
281
+ "#D6FF00": "dishwasher;dish;washer;dishwashing;machine",
282
+ "#00CCFF": "screen;silver;screen;projection;screen",
283
+ "#1400FF": "blanket;cover",
284
+ "#FFFF00": "sculpture",
285
+ "#0099FF": "hood;exhaust;hood",
286
+ "#0029FF": "sconce",
287
+ "#00FFCC": "vase",
288
+ "#2900FF": "traffic;light;traffic;signal;stoplight",
289
+ "#29FF00": "tray",
290
+ "#AD00FF": "ashcan;trash;can;garbage;can;wastebin;ash;bin;ash-bin;ashbin;dustbin;trash;barrel;trash;bin",
291
+ "#00F5FF": "fan",
292
+ "#4700FF": "pier;wharf;wharfage;dock",
293
+ "#7A00FF": "crt;screen",
294
+ "#00FFB8": "plate",
295
+ "#005CFF": "monitor;monitoring;device",
296
+ "#B8FF00": "bulletin;board;notice;board",
297
+ "#0085FF": "shower",
298
+ "#FFD600": "radiator",
299
+ "#19C2C2": "glass;drinking;glass",
300
+ "#66FF00": "clock",
301
+ "#5C00FF": "flag",
302
+ }
303
+
304
+
305
+ def ade_palette() -> List[List[int]]:
306
+ """ADE20K palette that maps each class to RGB values."""
307
+ return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
308
+ [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
309
+ [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
310
+ [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
311
+ [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
312
+ [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
313
+ [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
314
+ [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
315
+ [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
316
+ [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
317
+ [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
318
+ [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
319
+ [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
320
+ [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
321
+ [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
322
+ [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
323
+ [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
324
+ [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
325
+ [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
326
+ [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
327
+ [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
328
+ [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
329
+ [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
330
+ [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
331
+ [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
332
+ [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
333
+ [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
334
+ [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
335
+ [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
336
+ [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
337
+ [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
338
+ [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
339
+ [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
340
+ [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
341
+ [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
342
+ [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
343
+ [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
344
+ [102, 255, 0], [92, 0, 255]]
palette.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict
2
+ from colors import COLOR_MAPPING_, COLOR_MAPPING_CATEGORY_
3
+
4
+
5
+ def convert_hex_to_rgba(hex_code: str) -> str:
6
+ """Convert hex code to rgba.
7
+ Args:
8
+ hex_code (str): hex string
9
+ Returns:
10
+ str: rgba string
11
+ """
12
+ hex_code = hex_code.lstrip('#')
13
+ return "rgba(" + str(int(hex_code[0:2], 16)) + ", " + str(int(hex_code[2:4], 16)) + ", " + str(int(hex_code[4:6], 16)) + ", 1.0)"
14
+
15
+
16
+ def convert_dict_to_rgba(color_dict: Dict) -> Dict:
17
+ """Convert hex code to rgba for all elements in a dictionary.
18
+ Args:
19
+ color_dict (Dict): color dictionary
20
+ Returns:
21
+ Dict: color dictionary with rgba values
22
+ """
23
+ updated_dict = {}
24
+ for k, v in color_dict.items():
25
+ updated_dict[convert_hex_to_rgba(k)] = v
26
+ return updated_dict
27
+
28
+
29
+ def convert_nested_dict_to_rgba(nested_dict):
30
+ updated_dict = {}
31
+ for k, v in nested_dict.items():
32
+ updated_dict[k] = convert_dict_to_rgba(v)
33
+ return updated_dict
34
+
35
+
36
+ COLOR_MAPPING = convert_dict_to_rgba(COLOR_MAPPING_)
37
+ COLOR_MAPPING_CATEGORY = convert_nested_dict_to_rgba(COLOR_MAPPING_CATEGORY_)
utils.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gc
2
+
3
+ import numpy as np
4
+ from PIL import Image
5
+ import torch
6
+ from scipy.signal import fftconvolve
7
+
8
+ from palette import COLOR_MAPPING, COLOR_MAPPING_
9
+
10
+
11
+ def to_rgb(color: str) -> tuple:
12
+ """Convert hex color to rgb.
13
+ Args:
14
+ color (str): hex color
15
+ Returns:
16
+ tuple: rgb color
17
+ """
18
+ return tuple(int(color[i:i+2], 16) for i in (1, 3, 5))
19
+
20
+
21
+ def map_colors(color: str) -> str:
22
+ """Map color to hex value.
23
+ Args:
24
+ color (str): color name
25
+ Returns:
26
+ str: hex value
27
+ """
28
+ return COLOR_MAPPING[color]
29
+
30
+
31
+ def map_colors_rgb(color: tuple) -> str:
32
+ return COLOR_MAPPING_RGB[color]
33
+
34
+
35
+ def convolution(mask: Image.Image, size=9) -> Image:
36
+ """Method to blur the mask
37
+ Args:
38
+ mask (Image): masking image
39
+ size (int, optional): size of the blur. Defaults to 9.
40
+ Returns:
41
+ Image: blurred mask
42
+ """
43
+ mask = np.array(mask.convert("L"))
44
+ conv = np.ones((size, size)) / size**2
45
+ mask_blended = fftconvolve(mask, conv, 'same')
46
+ mask_blended = mask_blended.astype(np.uint8).copy()
47
+
48
+ border = size
49
+
50
+ # replace borders with original values
51
+ mask_blended[:border, :] = mask[:border, :]
52
+ mask_blended[-border:, :] = mask[-border:, :]
53
+ mask_blended[:, :border] = mask[:, :border]
54
+ mask_blended[:, -border:] = mask[:, -border:]
55
+
56
+ return Image.fromarray(mask_blended).convert("L")
57
+
58
+
59
+ def flush():
60
+ gc.collect()
61
+ torch.cuda.empty_cache()
62
+
63
+
64
+ def postprocess_image_masking(inpainted: Image, image: Image,
65
+ mask: Image) -> Image:
66
+ """Method to postprocess the inpainted image
67
+ Args:
68
+ inpainted (Image): inpainted image
69
+ image (Image): original image
70
+ mask (Image): mask
71
+ Returns:
72
+ Image: inpainted image
73
+ """
74
+ final_inpainted = Image.composite(inpainted.convert("RGBA"),
75
+ image.convert("RGBA"), mask)
76
+ return final_inpainted.convert("RGB")
77
+
78
+
79
+ COLOR_NAMES = list(COLOR_MAPPING.keys())
80
+ COLOR_RGB = [to_rgb(k) for k in COLOR_MAPPING_.keys()] + [(0, 0, 0),
81
+ (255, 255, 255)]
82
+ INVERSE_COLORS = {v: to_rgb(k) for k, v in COLOR_MAPPING_.items()}
83
+ COLOR_MAPPING_RGB = {to_rgb(k): v for k, v in COLOR_MAPPING_.items()}