File size: 4,413 Bytes
0985519 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
'''
The script to visualize
'''
import cv2
import torch
import os
import numpy as np
import colorsys
import seaborn as sns
platte = sns.color_palette("Spectral", 100, as_cmap=True) # doesn't work
from typing import Iterable, Tuple
import colorsys
import itertools
from fractions import Fraction
from pprint import pprint
######## The code to generate high-contrastive colors for visualization ##########
def zenos_dichotomy() -> Iterable[Fraction]:
"""
http://en.wikipedia.org/wiki/1/2_%2B_1/4_%2B_1/8_%2B_1/16_%2B_%C2%B7_%C2%B7_%C2%B7
"""
for k in itertools.count():
yield Fraction(1,2**k)
def fracs() -> Iterable[Fraction]:
"""
[Fraction(0, 1), Fraction(1, 2), Fraction(1, 4), Fraction(3, 4), Fraction(1, 8), Fraction(3, 8), Fraction(5, 8), Fraction(7, 8), Fraction(1, 16), Fraction(3, 16), ...]
[0.0, 0.5, 0.25, 0.75, 0.125, 0.375, 0.625, 0.875, 0.0625, 0.1875, ...]
"""
yield Fraction(0)
for k in zenos_dichotomy():
i = k.denominator # [1,2,4,8,16,...]
for j in range(1,i,2):
yield Fraction(j,i)
# can be used for the v in hsv to map linear values 0..1 to something that looks equidistant
# bias = lambda x: (math.sqrt(x/3)/Fraction(2,3)+Fraction(1,3))/Fraction(6,5)
HSVTuple = Tuple[Fraction, Fraction, Fraction]
RGBTuple = Tuple[float, float, float]
def hue_to_tones(h: Fraction) -> Iterable[HSVTuple]:
for s in [Fraction(6,10)]: # optionally use range
for v in [Fraction(8,10),Fraction(5,10)]: # could use range too
yield (h, s, v) # use bias for v here if you use range
def hsv_to_rgb(x: HSVTuple) -> RGBTuple:
return colorsys.hsv_to_rgb(*map(float, x))
flatten = itertools.chain.from_iterable
def hsvs() -> Iterable[HSVTuple]:
return flatten(map(hue_to_tones, fracs()))
def rgbs() -> Iterable[RGBTuple]:
return map(hsv_to_rgb, hsvs())
def rgb_to_css(x: RGBTuple) -> str:
uint8tuple = map(lambda y: int(y*255), x)
rgb_str = "{},{},{}".format(*uint8tuple)
rgb_value = rgb_str.split(",")
rgb_value = [int(d) for d in rgb_value]
return (rgb_value[0], rgb_value[1], rgb_value[2])
def css_colors() -> Iterable[str]:
return map(rgb_to_css, rgbs())
sample_colors = list(itertools.islice(css_colors(), 400))
########## ########## ########## ########## ########## ########## ##########
def draw_box(im, box, thickness=1, color=(255,0,0), trackids=[]):
x1, y1, w, h = box
x2 = x1+w
y2 = y1+h
cv2.rectangle(im, (x1, y1), (x2, y2), color=color, thickness=thickness)
trackids = [str(id) for id in trackids]
cv2.putText(im, ",".join(trackids), (int(x1+2), int(y1+12)), cv2.FONT_HERSHEY_COMPLEX, 0.6, color, 2)
return im
def draw_pieces(pieces_annos, img_dir, save_dir):
os.makedirs(save_dir, exist_ok=True)
frames = torch.unique(pieces_annos[:,:,0])
for frame in frames:
frame_bboxes = pieces_annos[pieces_annos[:, :, 0]==frame]
img = os.path.join(img_dir, "%06d.jpg" % frame)
im = cv2.imread(img)
for bbox in frame_bboxes:
trackid = bbox[1]
if trackid == 0:
continue
else:
coord = bbox[2:6]
occupy_indices = (frame_bboxes[:, 2:6] == coord)[:,0]
bboxes = frame_bboxes[occupy_indices]
trackids = torch.unique(bboxes[:, 1]).int().tolist()
im = draw_box(im, coord.numpy(), thickness=2, color=sample_colors[int(trackid)], trackids=trackids)
save_path = os.path.join(save_dir, "vis_%06d.jpg" % frame)
cv2.imwrite(save_path, im)
if __name__ == "__main__":
seq_name = "MOT17-02-DPM"
img_dir = "data/MOT17/train/{}/img1".format(seq_name)
anno = "data/MOT17/train_pieces/{}".format(seq_name)
index = 2895
piece_anno = torch.load("{}/{}_pieces.pth".format(anno, index))
# import pdb; pdb.set_trace()
for i in range(piece_anno.shape[0]):
piece_anno[i, piece_anno[i, :, 1] !=0, 1] = i+1
tracklet_anno = torch.load("{}/{}_tracklets.pth".format(anno, index))
draw_pieces(piece_anno, img_dir, "visualizations/{}/{}".format(seq_name, index))
cmd = "ffmpeg -y -r 10 -i visualizations/{}/{}/vis_%06d.jpg vis_{}_{}.mp4".format(seq_name, index, seq_name, index)
os.system(cmd)
|