Spaces:
Runtime error
Runtime error
# -*- coding: utf-8 -*- | |
""" | |
@File : visualizer.py | |
@Time : 2022/04/05 11:39:33 | |
@Author : Shilong Liu | |
@Contact : slongliu86@gmail.com | |
""" | |
import datetime | |
import os | |
import cv2 | |
import matplotlib.pyplot as plt | |
import numpy as np | |
import torch | |
from matplotlib import transforms | |
from matplotlib.collections import PatchCollection | |
from matplotlib.patches import Polygon | |
from pycocotools import mask as maskUtils | |
def renorm( | |
img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] | |
) -> torch.FloatTensor: | |
# img: tensor(3,H,W) or tensor(B,3,H,W) | |
# return: same as img | |
assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim() | |
if img.dim() == 3: | |
assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % ( | |
img.size(0), | |
str(img.size()), | |
) | |
img_perm = img.permute(1, 2, 0) | |
mean = torch.Tensor(mean) | |
std = torch.Tensor(std) | |
img_res = img_perm * std + mean | |
return img_res.permute(2, 0, 1) | |
else: # img.dim() == 4 | |
assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % ( | |
img.size(1), | |
str(img.size()), | |
) | |
img_perm = img.permute(0, 2, 3, 1) | |
mean = torch.Tensor(mean) | |
std = torch.Tensor(std) | |
img_res = img_perm * std + mean | |
return img_res.permute(0, 3, 1, 2) | |
class ColorMap: | |
def __init__(self, basergb=[255, 255, 0]): | |
self.basergb = np.array(basergb) | |
def __call__(self, attnmap): | |
# attnmap: h, w. np.uint8. | |
# return: h, w, 4. np.uint8. | |
assert attnmap.dtype == np.uint8 | |
h, w = attnmap.shape | |
res = self.basergb.copy() | |
res = res[None][None].repeat(h, 0).repeat(w, 1) # h, w, 3 | |
attn1 = attnmap.copy()[..., None] # h, w, 1 | |
res = np.concatenate((res, attn1), axis=-1).astype(np.uint8) | |
return res | |
def rainbow_text(x, y, ls, lc, **kw): | |
""" | |
Take a list of strings ``ls`` and colors ``lc`` and place them next to each | |
other, with text ls[i] being shown in color lc[i]. | |
This example shows how to do both vertical and horizontal text, and will | |
pass all keyword arguments to plt.text, so you can set the font size, | |
family, etc. | |
""" | |
t = plt.gca().transData | |
fig = plt.gcf() | |
plt.show() | |
# horizontal version | |
for s, c in zip(ls, lc): | |
text = plt.text(x, y, " " + s + " ", color=c, transform=t, **kw) | |
text.draw(fig.canvas.get_renderer()) | |
ex = text.get_window_extent() | |
t = transforms.offset_copy(text._transform, x=ex.width, units="dots") | |
# #vertical version | |
# for s,c in zip(ls,lc): | |
# text = plt.text(x,y," "+s+" ",color=c, transform=t, | |
# rotation=90,va='bottom',ha='center',**kw) | |
# text.draw(fig.canvas.get_renderer()) | |
# ex = text.get_window_extent() | |
# t = transforms.offset_copy(text._transform, y=ex.height, units='dots') | |
class COCOVisualizer: | |
def __init__(self, coco=None, tokenlizer=None) -> None: | |
self.coco = coco | |
def visualize(self, img, tgt, caption=None, dpi=180, savedir="vis"): | |
""" | |
img: tensor(3, H, W) | |
tgt: make sure they are all on cpu. | |
must have items: 'image_id', 'boxes', 'size' | |
""" | |
plt.figure(dpi=dpi) | |
plt.rcParams["font.size"] = "5" | |
ax = plt.gca() | |
img = renorm(img).permute(1, 2, 0) | |
# if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': | |
# import ipdb; ipdb.set_trace() | |
ax.imshow(img) | |
self.addtgt(tgt) | |
if tgt is None: | |
image_id = 0 | |
elif "image_id" not in tgt: | |
image_id = 0 | |
else: | |
image_id = tgt["image_id"] | |
if caption is None: | |
savename = "{}/{}-{}.png".format( | |
savedir, int(image_id), str(datetime.datetime.now()).replace(" ", "-") | |
) | |
else: | |
savename = "{}/{}-{}-{}.png".format( | |
savedir, caption, int(image_id), str(datetime.datetime.now()).replace(" ", "-") | |
) | |
print("savename: {}".format(savename)) | |
os.makedirs(os.path.dirname(savename), exist_ok=True) | |
plt.savefig(savename) | |
plt.close() | |
def addtgt(self, tgt): | |
""" """ | |
if tgt is None or not "boxes" in tgt: | |
ax = plt.gca() | |
if "caption" in tgt: | |
ax.set_title(tgt["caption"], wrap=True) | |
ax.set_axis_off() | |
return | |
ax = plt.gca() | |
H, W = tgt["size"] | |
numbox = tgt["boxes"].shape[0] | |
color = [] | |
polygons = [] | |
boxes = [] | |
for box in tgt["boxes"].cpu(): | |
unnormbbox = box * torch.Tensor([W, H, W, H]) | |
unnormbbox[:2] -= unnormbbox[2:] / 2 | |
[bbox_x, bbox_y, bbox_w, bbox_h] = unnormbbox.tolist() | |
boxes.append([bbox_x, bbox_y, bbox_w, bbox_h]) | |
poly = [ | |
[bbox_x, bbox_y], | |
[bbox_x, bbox_y + bbox_h], | |
[bbox_x + bbox_w, bbox_y + bbox_h], | |
[bbox_x + bbox_w, bbox_y], | |
] | |
np_poly = np.array(poly).reshape((4, 2)) | |
polygons.append(Polygon(np_poly)) | |
c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0] | |
color.append(c) | |
p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.1) | |
ax.add_collection(p) | |
p = PatchCollection(polygons, facecolor="none", edgecolors=color, linewidths=2) | |
ax.add_collection(p) | |
if "strings_positive" in tgt and len(tgt["strings_positive"]) > 0: | |
assert ( | |
len(tgt["strings_positive"]) == numbox | |
), f"{len(tgt['strings_positive'])} = {numbox}, " | |
for idx, strlist in enumerate(tgt["strings_positive"]): | |
cate_id = int(tgt["labels"][idx]) | |
_string = str(cate_id) + ":" + " ".join(strlist) | |
bbox_x, bbox_y, bbox_w, bbox_h = boxes[idx] | |
# ax.text(bbox_x, bbox_y, _string, color='black', bbox={'facecolor': 'yellow', 'alpha': 1.0, 'pad': 1}) | |
ax.text( | |
bbox_x, | |
bbox_y, | |
_string, | |
color="black", | |
bbox={"facecolor": color[idx], "alpha": 0.6, "pad": 1}, | |
) | |
if "box_label" in tgt: | |
assert len(tgt["box_label"]) == numbox, f"{len(tgt['box_label'])} = {numbox}, " | |
for idx, bl in enumerate(tgt["box_label"]): | |
_string = str(bl) | |
bbox_x, bbox_y, bbox_w, bbox_h = boxes[idx] | |
# ax.text(bbox_x, bbox_y, _string, color='black', bbox={'facecolor': 'yellow', 'alpha': 1.0, 'pad': 1}) | |
ax.text( | |
bbox_x, | |
bbox_y, | |
_string, | |
color="black", | |
bbox={"facecolor": color[idx], "alpha": 0.6, "pad": 1}, | |
) | |
if "caption" in tgt: | |
ax.set_title(tgt["caption"], wrap=True) | |
# plt.figure() | |
# rainbow_text(0.0,0.0,"all unicorns poop rainbows ! ! !".split(), | |
# ['red', 'orange', 'brown', 'green', 'blue', 'purple', 'black']) | |
if "attn" in tgt: | |
# if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': | |
# import ipdb; ipdb.set_trace() | |
if isinstance(tgt["attn"], tuple): | |
tgt["attn"] = [tgt["attn"]] | |
for item in tgt["attn"]: | |
attn_map, basergb = item | |
attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-3) | |
attn_map = (attn_map * 255).astype(np.uint8) | |
cm = ColorMap(basergb) | |
heatmap = cm(attn_map) | |
ax.imshow(heatmap) | |
ax.set_axis_off() | |
def showAnns(self, anns, draw_bbox=False): | |
""" | |
Display the specified annotations. | |
:param anns (array of object): annotations to display | |
:return: None | |
""" | |
if len(anns) == 0: | |
return 0 | |
if "segmentation" in anns[0] or "keypoints" in anns[0]: | |
datasetType = "instances" | |
elif "caption" in anns[0]: | |
datasetType = "captions" | |
else: | |
raise Exception("datasetType not supported") | |
if datasetType == "instances": | |
ax = plt.gca() | |
ax.set_autoscale_on(False) | |
polygons = [] | |
color = [] | |
for ann in anns: | |
c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0] | |
if "segmentation" in ann: | |
if type(ann["segmentation"]) == list: | |
# polygon | |
for seg in ann["segmentation"]: | |
poly = np.array(seg).reshape((int(len(seg) / 2), 2)) | |
polygons.append(Polygon(poly)) | |
color.append(c) | |
else: | |
# mask | |
t = self.imgs[ann["image_id"]] | |
if type(ann["segmentation"]["counts"]) == list: | |
rle = maskUtils.frPyObjects( | |
[ann["segmentation"]], t["height"], t["width"] | |
) | |
else: | |
rle = [ann["segmentation"]] | |
m = maskUtils.decode(rle) | |
img = np.ones((m.shape[0], m.shape[1], 3)) | |
if ann["iscrowd"] == 1: | |
color_mask = np.array([2.0, 166.0, 101.0]) / 255 | |
if ann["iscrowd"] == 0: | |
color_mask = np.random.random((1, 3)).tolist()[0] | |
for i in range(3): | |
img[:, :, i] = color_mask[i] | |
ax.imshow(np.dstack((img, m * 0.5))) | |
if "keypoints" in ann and type(ann["keypoints"]) == list: | |
# turn skeleton into zero-based index | |
sks = np.array(self.loadCats(ann["category_id"])[0]["skeleton"]) - 1 | |
kp = np.array(ann["keypoints"]) | |
x = kp[0::3] | |
y = kp[1::3] | |
v = kp[2::3] | |
for sk in sks: | |
if np.all(v[sk] > 0): | |
plt.plot(x[sk], y[sk], linewidth=3, color=c) | |
plt.plot( | |
x[v > 0], | |
y[v > 0], | |
"o", | |
markersize=8, | |
markerfacecolor=c, | |
markeredgecolor="k", | |
markeredgewidth=2, | |
) | |
plt.plot( | |
x[v > 1], | |
y[v > 1], | |
"o", | |
markersize=8, | |
markerfacecolor=c, | |
markeredgecolor=c, | |
markeredgewidth=2, | |
) | |
if draw_bbox: | |
[bbox_x, bbox_y, bbox_w, bbox_h] = ann["bbox"] | |
poly = [ | |
[bbox_x, bbox_y], | |
[bbox_x, bbox_y + bbox_h], | |
[bbox_x + bbox_w, bbox_y + bbox_h], | |
[bbox_x + bbox_w, bbox_y], | |
] | |
np_poly = np.array(poly).reshape((4, 2)) | |
polygons.append(Polygon(np_poly)) | |
color.append(c) | |
# p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.4) | |
# ax.add_collection(p) | |
p = PatchCollection(polygons, facecolor="none", edgecolors=color, linewidths=2) | |
ax.add_collection(p) | |
elif datasetType == "captions": | |
for ann in anns: | |
print(ann["caption"]) | |