Apply Turbo colormap to a normalized value tensor. For cuda acceleration, the colormap is preloaded as a tensor. Args: value (torch.Tensor): A tensor with values in the range [0, 1]. Returns: torch.Tensor: A tensor with the same shape as input, but with an addi
(value)
| 171 | |
| 172 | turbo_colormap = torch.tensor(matplotlib.colormaps.get_cmap('turbo').colors, device='cuda') |
| 173 | def cmap(value): |
| 174 | """ |
| 175 | Apply Turbo colormap to a normalized value tensor. For cuda acceleration, the colormap is preloaded as a tensor. |
| 176 | |
| 177 | Args: |
| 178 | value (torch.Tensor): A tensor with values in the range [0, 1]. |
| 179 | |
| 180 | Returns: |
| 181 | torch.Tensor: A tensor with the same shape as input, but with an additional |
| 182 | dimension for RGB channels. |
| 183 | """ |
| 184 | # assert torch.all(value >= 0) and torch.all(value <= 1), "Input values should be in the range [0, 1]" |
| 185 | |
| 186 | # Scale the input values to the range [0, len(TURBO_COLORS) - 1] |
| 187 | indices = (value * (turbo_colormap.shape[0] - 1)).long() |
| 188 | # Gather the corresponding colors |
| 189 | colored = turbo_colormap[indices] |
| 190 | return colored |
| 191 | |
| 192 | def clip_color(cos_sim, bg_mask, height, width, thresh=0.7, res_finetuned=False, coloring=False, device='cuda'): |
| 193 | # 着色方案不一样 |