kahnchana commited on
Commit
1d7cddb
1 Parent(s): ed295cc
Files changed (5) hide show
  1. .gitignore +135 -0
  2. app.py +67 -0
  3. arial.ttf +0 -0
  4. infer_model.py +202 -0
  5. utils.py +166 -0
.gitignore ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # MacOS & PyCharm
2
+ .DS_store
3
+ .DS_Store
4
+ .idea/
5
+ sync.sh
6
+
7
+ # Byte-compiled / optimized / DLL files
8
+ __pycache__/
9
+ *.py[cod]
10
+ *$py.class
11
+
12
+ # C extensions
13
+ *.so
14
+
15
+ # Distribution / packaging
16
+ .Python
17
+ build/
18
+ develop-eggs/
19
+ dist/
20
+ downloads/
21
+ eggs/
22
+ .eggs/
23
+ lib/
24
+ lib64/
25
+ parts/
26
+ sdist/
27
+ var/
28
+ wheels/
29
+ pip-wheel-metadata/
30
+ share/python-wheels/
31
+ *.egg-info/
32
+ .installed.cfg
33
+ *.egg
34
+ MANIFEST
35
+
36
+ # PyInstaller
37
+ # Usually these files are written by a python script from a template
38
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
39
+ *.manifest
40
+ *.spec
41
+
42
+ # Installer logs
43
+ pip-log.txt
44
+ pip-delete-this-directory.txt
45
+
46
+ # Unit test / coverage reports
47
+ htmlcov/
48
+ .tox/
49
+ .nox/
50
+ .coverage
51
+ .coverage.*
52
+ .cache
53
+ nosetests.xml
54
+ coverage.xml
55
+ *.cover
56
+ *.py,cover
57
+ .hypothesis/
58
+ .pytest_cache/
59
+
60
+ # Translations
61
+ *.mo
62
+ *.pot
63
+
64
+ # Django stuff:
65
+ *.log
66
+ local_settings.py
67
+ db.sqlite3
68
+ db.sqlite3-journal
69
+
70
+ # Flask stuff:
71
+ instance/
72
+ .webassets-cache
73
+
74
+ # Scrapy stuff:
75
+ .scrapy
76
+
77
+ # Sphinx documentation
78
+ docs/_build/
79
+
80
+ # PyBuilder
81
+ target/
82
+
83
+ # Jupyter Notebook
84
+ .ipynb_checkpoints
85
+
86
+ # IPython
87
+ profile_default/
88
+ ipython_config.py
89
+
90
+ # pyenv
91
+ .python-version
92
+
93
+ # pipenv
94
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
95
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
96
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
97
+ # install all needed dependencies.
98
+ #Pipfile.lock
99
+
100
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow
101
+ __pypackages__/
102
+
103
+ # Celery stuff
104
+ celerybeat-schedule
105
+ celerybeat.pid
106
+
107
+ # SageMath parsed files
108
+ *.sage.py
109
+
110
+ # Environments
111
+ .env
112
+ .venv
113
+ env/
114
+ venv/
115
+ ENV/
116
+ env.bak/
117
+ venv.bak/
118
+
119
+ # Spyder project settings
120
+ .spyderproject
121
+ .spyproject
122
+
123
+ # Rope project settings
124
+ .ropeproject
125
+
126
+ # mkdocs documentation
127
+ /site
128
+
129
+ # mypy
130
+ .mypy_cache/
131
+ .dmypy.json
132
+ dmypy.json
133
+
134
+ # Pyre type checker
135
+ .pyre/
app.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ import torch
4
+ from PIL import Image
5
+
6
+ from infer_model import CLIPpyModel
7
+ from utils import get_similarity, get_transform, ade_palette, get_cmap_image
8
+
9
+ pretrained_ckpt = "https://github.com/kahnchana/clippy/releases/download/v1.0/clippy_5k.pt"
10
+ ckpt = torch.utils.model_zoo.load_url(pretrained_ckpt)
11
+
12
+ clippy = CLIPpyModel()
13
+ transform = get_transform((224, 224))
14
+
15
+ msg = clippy.load_state_dict(ckpt, strict=False)
16
+
17
+ palette = ade_palette()
18
+
19
+
20
+ def process_image(img, captions):
21
+ sample_text = [x.strip() for x in captions.split(",")]
22
+ sample_prompts = [f"a photo of a {x}" for x in sample_text]
23
+
24
+ image = Image.fromarray(img)
25
+ image_vector = clippy.encode_image(transform(image).unsqueeze(0), get_pos_tokens=True)
26
+ text_vector = clippy.text.encode(sample_prompts, convert_to_tensor=True)
27
+
28
+ similarity = get_similarity(image_vector, text_vector, (224, 224), do_argmax=True)[0, 0].numpy()
29
+ rgb_seg = np.zeros((similarity.shape[0], similarity.shape[1], 3), dtype=np.uint8)
30
+ for idx, _ in enumerate(sample_text):
31
+ rgb_seg[similarity == idx] = palette[idx]
32
+
33
+ joint = Image.blend(image, Image.fromarray(rgb_seg), 0.5)
34
+ cmap = get_cmap_image({label: tuple(palette[idx]) for idx, label in enumerate(sample_text)})
35
+
36
+ return cmap, rgb_seg, joint
37
+
38
+
39
+ title = 'CLIPpy'
40
+
41
+ description = """
42
+ Gradio Demo for CLIPpy: Perceptual Grouping in Contrastive Vision Language Models. \n \n
43
+ Upload an image and type in a set of comma separated labels (e.g.: "man, woman, background").
44
+ CLIPPy will segment the image, according to the set of class label you provide.
45
+ """
46
+
47
+ article = """
48
+ <p style='text-align: center'>
49
+ <a href='https://arxiv.org/abs/2210.09996' target='_blank'>
50
+ Perceptual Grouping in Contrastive Vision Language Models
51
+ </a>
52
+ |
53
+ <a href='https://github.com/kahnchana/clippy' target='_blank'>Github Repository</a></p>
54
+ """
55
+
56
+ demo = gr.Interface(
57
+ fn=process_image,
58
+ inputs=[gr.Image(shape=(224, 224)), "text"],
59
+ outputs=[gr.Image(shape=(224, 224)).style(height=150),
60
+ gr.Image(shape=(224, 224)).style(height=260),
61
+ gr.Image(shape=(224, 224)).style(height=260)],
62
+ title=title,
63
+ description=description,
64
+ article=article,
65
+ )
66
+
67
+ demo.launch()
arial.ttf ADDED
Binary file (367 kB). View file
 
infer_model.py ADDED
@@ -0,0 +1,202 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import abc
2
+ import math
3
+
4
+ import torch
5
+ import torch.nn.functional as F
6
+ from sentence_transformers import SentenceTransformer
7
+ from timm.models.vision_transformer import (
8
+ VisionTransformer,
9
+ build_model_with_cfg,
10
+ checkpoint_filter_fn,
11
+ checkpoint_seq,
12
+ resolve_pretrained_cfg,
13
+ )
14
+ from torch import Tensor, nn
15
+
16
+
17
+ class BlankLayer(nn.Module):
18
+ pass
19
+
20
+
21
+ class CustomViT(VisionTransformer):
22
+ def __init__(
23
+ self,
24
+ *args,
25
+ image_pooling="gmp",
26
+ **kwargs,
27
+ ):
28
+ super(CustomViT, self).__init__(
29
+ *args, **kwargs
30
+ )
31
+ self.image_pooling = image_pooling
32
+
33
+ def forward_head(self, x, pre_logits: bool = False):
34
+ if self.image_pooling:
35
+ if self.image_pooling == "gap":
36
+ x = x[:, self.num_prefix_tokens:].mean(dim=1)
37
+ elif self.image_pooling == "gmp":
38
+ x = x[:, self.num_prefix_tokens:].max(dim=-2)[0]
39
+ elif self.image_pooling == "all":
40
+ x = x[:, self.num_prefix_tokens:]
41
+ else: # cls by default
42
+ x = x[:, 0]
43
+ x = self.fc_norm(x)
44
+ return x if pre_logits else self.head(x)
45
+
46
+ def forward(self, x, get_pos_tokens=False):
47
+ x = self.forward_features(x, get_pos_tokens=get_pos_tokens)
48
+ if get_pos_tokens:
49
+ return self.fc_norm(x[:, self.num_prefix_tokens:])
50
+ x = self.forward_head(x)
51
+ return x
52
+
53
+ def forward_features(self, x, get_pos_tokens=False):
54
+ _, nc, h, w = x.shape
55
+ x = self.patch_embed(x)
56
+ x = self._pos_embed(x, w, h)
57
+ if self.grad_checkpointing and not torch.jit.is_scripting():
58
+ x = checkpoint_seq(self.blocks, x)
59
+ else:
60
+ x = self.blocks(x)
61
+ x = self.norm(x)
62
+ return x
63
+
64
+ def _pos_embed(self, x, w, h):
65
+ if self.no_embed_class:
66
+ # deit-3, updated JAX (big vision)
67
+ # position embedding does not overlap with class token, add then concat
68
+ x = x + self.pos_embed
69
+ if self.cls_token is not None:
70
+ x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
71
+ else:
72
+ # original timm, JAX, and deit vit impl
73
+ # pos_embed has entry for class token, concat then add
74
+ if self.cls_token is not None:
75
+ x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
76
+ x = x + self._interpolate_pos_encoding(x, w, h)
77
+ return self.pos_drop(x)
78
+
79
+ def _interpolate_pos_encoding(self, x, w, h):
80
+ npatch = x.shape[1] - 1
81
+ N = self.pos_embed.shape[1] - 1
82
+ if npatch == N and w == h:
83
+ return self.pos_embed
84
+ class_pos_embed = self.pos_embed[:, 0]
85
+ patch_pos_embed = self.pos_embed[:, 1:]
86
+ dim = x.shape[-1]
87
+ w0 = w // self.patch_embed.patch_size[0]
88
+ h0 = h // self.patch_embed.patch_size[1]
89
+ # we add a small number to avoid floating point error in the interpolation
90
+ # see discussion at https://github.com/facebookresearch/dino/issues/8
91
+ w0, h0 = w0 + 0.1, h0 + 0.1
92
+ patch_pos_embed = nn.functional.interpolate(
93
+ patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(
94
+ 0, 3, 1, 2
95
+ ),
96
+ scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
97
+ mode="bicubic",
98
+ )
99
+ assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
100
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
101
+ return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
102
+
103
+
104
+ def _create_vision_transformer(variant, pretrained=False, **kwargs):
105
+ if kwargs.get("features_only", None):
106
+ raise RuntimeError("features_only not implemented for Vision Transformer models.")
107
+
108
+ pretrained_cfg = resolve_pretrained_cfg(
109
+ variant, pretrained_cfg=kwargs.pop("pretrained_cfg", None)
110
+ )
111
+ model = build_model_with_cfg(
112
+ CustomViT,
113
+ variant,
114
+ pretrained,
115
+ pretrained_cfg=pretrained_cfg,
116
+ pretrained_filter_fn=checkpoint_filter_fn,
117
+ pretrained_custom_load="npz" in pretrained_cfg["url"],
118
+ **kwargs,
119
+ )
120
+ return model
121
+
122
+
123
+ def vit_base_patch16_224(pretrained=False, variant="vit_base_patch16_224_dino", **kwargs):
124
+ """ViT-Base (ViT-B/16) /w DINO pretrained weights (no head) - https://arxiv.org/abs/2104.14294"""
125
+ model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
126
+ model = _create_vision_transformer(variant, pretrained=pretrained, **model_kwargs)
127
+ return model
128
+
129
+
130
+ class CLIPpyModel(abc.ABC, torch.nn.Module):
131
+ """ Implements code for running inference with pre-trained CLIPpy model.
132
+
133
+ NOTE: weights used are for a model trained with lower batch-size leading to results below those in paper.
134
+ """
135
+
136
+ def __init__(
137
+ self,
138
+ image_pooling: str = "cls",
139
+ text_pooling: str = "gap",
140
+ ):
141
+ super().__init__()
142
+
143
+ self.visual = BlankLayer()
144
+
145
+ self.visual.trunk = vit_base_patch16_224(True, image_pooling=image_pooling)
146
+
147
+ self.text = SentenceTransformer("sentence-transformers/sentence-t5-base")
148
+ self.logit_scale = nn.Parameter(torch.ones([]) * math.log(1 / 0.07))
149
+ self.set_text_pooling(text_pooling)
150
+
151
+ self._divisor_eps = 1e-4
152
+ self._image_pooling = image_pooling
153
+ self._text_pooling = text_pooling
154
+
155
+ def forward(
156
+ self,
157
+ images: Tensor,
158
+ input_ids: Tensor,
159
+ input_id_masks: Tensor,
160
+ get_pos_tokens: bool = False,
161
+ **kwargs,
162
+ ):
163
+
164
+ image_encodings = self.encode_image(images, get_pos_tokens=get_pos_tokens)
165
+
166
+ if get_pos_tokens:
167
+ return {
168
+ image_encodings: image_encodings,
169
+ }
170
+
171
+ text_encodings = self.encode_text(input_ids, input_id_masks)
172
+
173
+ return {
174
+ image_encodings: image_encodings,
175
+ text_encodings: text_encodings,
176
+ }
177
+
178
+ def encode_text(self, input_ids: Tensor, input_id_masks: Tensor = None, **kwargs):
179
+ output = self.text({"input_ids": input_ids, "attention_mask": input_id_masks})[
180
+ "sentence_embedding"
181
+ ]
182
+ return self.text_head(output)
183
+
184
+ def text_head(self, hidden_states: Tensor, input_id_masks: Tensor = None, **kwargs):
185
+ return F.normalize(hidden_states, dim=-1, eps=self._divisor_eps).float()
186
+
187
+ def encode_image(self, images: Tensor, get_pos_tokens: bool = False, **kwargs):
188
+ output = self.visual.trunk(images, get_pos_tokens)
189
+ return self.image_head(output, get_pos_tokens=get_pos_tokens)
190
+
191
+ def image_head(self, hidden_states: Tensor, get_pos_tokens: bool = False, **kwargs):
192
+ return F.normalize(hidden_states, dim=-1, eps=self._divisor_eps).float()
193
+
194
+ def set_text_pooling(self, pooling):
195
+ """ Converts pooling in the Hugging Face model to be max or average pooling"""
196
+ if pooling == "gmp":
197
+ self.text[1].pooling_mode_mean_tokens = False
198
+ self.text[1].pooling_mode_max_tokens = True
199
+ elif pooling == "gap":
200
+ pass
201
+ else:
202
+ raise NotImplementedError(f"{pooling} not implemented")
utils.py ADDED
@@ -0,0 +1,166 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import matplotlib
2
+ import matplotlib.cm as cm
3
+ import matplotlib.colors as mcolors
4
+ import numpy as np
5
+ import torch
6
+ import torchvision
7
+ from PIL import Image, ImageDraw, ImageFont
8
+ from einops import rearrange
9
+ from matplotlib import pyplot as plt
10
+
11
+
12
+ def get_similarity(image_encodings, label_encodings, target_shape, interpolation="bilinear", do_argmax=False):
13
+ """
14
+
15
+ Args:
16
+ image_encodings:
17
+ label_encodings:
18
+ target_shape:
19
+ interpolation: nearest, bilinear
20
+ do_argmax:
21
+
22
+ Returns:
23
+
24
+ """
25
+
26
+ image_encodings = image_encodings.cpu()
27
+ label_encodings = label_encodings.cpu()
28
+
29
+ image_encodings = rearrange(
30
+ image_encodings, "b (h w) d -> d b h w", h=int(np.sqrt(image_encodings.shape[-2]))
31
+ )
32
+ # assuming square inputs & targets
33
+ scale_ratio = (target_shape[-2] / image_encodings.shape[-2],
34
+ target_shape[-1] / image_encodings.shape[-1],)
35
+ temp_list = []
36
+ for i in image_encodings:
37
+ i = i.unsqueeze(1)
38
+ i = torch.nn.functional.interpolate(
39
+ i, scale_factor=scale_ratio, mode=interpolation
40
+ )
41
+ temp_list.append(i)
42
+ image_encodings = torch.cat(temp_list, dim=1)
43
+
44
+ image_encodings = rearrange(image_encodings, "b d h w -> b h w d")
45
+ similarity = image_encodings @ label_encodings.T
46
+ similarity = rearrange(similarity, "b h w d-> b d h w")
47
+ if do_argmax:
48
+ similarity = torch.argmax(similarity, dim=1, keepdim=True).to(torch.float64)
49
+ return similarity
50
+
51
+
52
+ def get_cmap(ncolors):
53
+ if ncolors > 9:
54
+ cmap = plt.cm.tab20
55
+ else:
56
+ cmap = plt.cm.tab10
57
+ cmaplist = [cmap(i) for i in range(ncolors)]
58
+ cmap = matplotlib.colors.LinearSegmentedColormap.from_list("custom", cmaplist, ncolors)
59
+
60
+ mappable = cm.ScalarMappable(cmap=cmap)
61
+ mappable.set_array([])
62
+ mappable.set_clim(-0.5, ncolors + 0.5)
63
+
64
+ return cmap, mappable
65
+
66
+
67
+ def vis_prediction(sample_text, img_arr, similarity):
68
+ N = len(sample_text)
69
+ cmap, mappable = get_cmap(N)
70
+
71
+ fig, axs = plt.subplots(1, 2)
72
+
73
+ _ = axs[0].imshow(img_arr)
74
+ _ = axs[1].imshow(img_arr)
75
+ _ = axs[1].imshow(similarity, cmap=cmap, interpolation="nearest", vmin=0, vmax=N, alpha=0.5)
76
+ axs[0].axis("off")
77
+ axs[1].axis("off")
78
+
79
+ fig.subplots_adjust(bottom=0.2)
80
+ cbar_ax = fig.add_axes([0.0, 0.85, 1.0, 0.05])
81
+ colorbar = plt.colorbar(mappable, cax=cbar_ax, cmap=cmap, orientation="horizontal")
82
+ colorbar.set_ticks(np.linspace(0, N, N))
83
+ colorbar.set_ticklabels(sample_text)
84
+
85
+ return fig
86
+
87
+
88
+ class DummyArgs:
89
+ def __init__(self, **kwargs):
90
+ self.__dict__.update(kwargs)
91
+
92
+
93
+ def get_transform(size=(224, 224)):
94
+ transform = torchvision.transforms.Compose([
95
+ torchvision.transforms.Resize(size),
96
+ torchvision.transforms.ToTensor(),
97
+ torchvision.transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073),
98
+ std=(0.26862954, 0.26130258, 0.27577711))
99
+ ])
100
+ return transform
101
+
102
+
103
+ def ade_palette():
104
+ """ADE20K palette that maps each class to RGB values."""
105
+ return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
106
+ [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
107
+ [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
108
+ [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
109
+ [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
110
+ [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
111
+ [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
112
+ [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
113
+ [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
114
+ [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
115
+ [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
116
+ [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
117
+ [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
118
+ [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
119
+ [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
120
+ [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
121
+ [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
122
+ [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
123
+ [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
124
+ [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
125
+ [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
126
+ [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
127
+ [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
128
+ [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
129
+ [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
130
+ [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
131
+ [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
132
+ [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
133
+ [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
134
+ [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
135
+ [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
136
+ [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
137
+ [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
138
+ [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
139
+ [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
140
+ [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
141
+ [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
142
+ [102, 255, 0], [92, 0, 255]]
143
+
144
+
145
+ def get_cmap_image(legend):
146
+ # Define the size of the legend image
147
+ width = 200
148
+ height = len(legend) * 20
149
+
150
+ # Create a new image with the desired size and background color
151
+ img = Image.new('RGB', (width, height), (255, 255, 255))
152
+
153
+ # Create a drawing context
154
+ draw = ImageDraw.Draw(img)
155
+
156
+ # Define the font to use for the legend labels
157
+ font = ImageFont.truetype('arial.ttf', 16)
158
+
159
+ # Loop through the items in legend and draw a rectangle and label for each
160
+ y = 0
161
+ for label, color in legend.items():
162
+ draw.rectangle((0, y, 20, y + 20), fill=color)
163
+ draw.text((30, y), label, font=font, fill=(0, 0, 0))
164
+ y += 20
165
+
166
+ return img