def __init__(self,
cfg,
confidence_threshold=-1,
show_mask_heatmaps=False,
show_mask_montage=True,
masks_per_dim=2,
categories=None):
self.categories = categories
self.cfg = cfg.clone()
mask_threshold = -1 if show_mask_montage else 0.5
self.masker = Masker(threshold=mask_threshold, padding=1)
if confidence_threshold >= 0:
self.confidence_threshold = confidence_threshold
else:
self.confidence_threshold = cfg.MODEL.ROI_HEADS.SCORE_THRESH_VISUALIZATION
self.show_heatmap = show_mask_heatmaps
self.show_mask_montage = show_mask_montage
self.masks_per_dim = masks_per_dim
# used to make colors for each class
self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1])
self.my_palette = []
colors = create_palette()
self.my_palette = []
for color in colors:
self.my_palette.append(color['rgb'])
def visualize_tracking_object_class(dir_image,
dir_json,
dir_output=None,
min_threshold=0,
max_threshold=1,
thick=5,
colors=None,
display_caption=True,
CATEGORIES=[]):
if colors is None:
palette = create_palette()
colors = {}
for x in palette:
colors[x['id']] = x['rgb']
image = cv2.imread(dir_image)
if os.path.exists(dir_json):
with open(dir_json) as f:
data = json.load(f)
if isinstance(data, dict):
data = [x for x in data.values()]
for bbox in data:
if 'score' not in bbox.keys() or (bbox['score'] >= min_threshold
and
bbox['score'] <= max_threshold):
top_left = (int(bbox['bbox']['x']), int(bbox['bbox']['y']))
bottom_right = (int(bbox['bbox']['x']) +
int(bbox['bbox']['w']),
int(bbox['bbox']['y']) +
int(bbox['bbox']['h']))
image = cv2.rectangle(image, tuple(top_left),
tuple(bottom_right),
tuple(colors[bbox['category_id']]),
thick)
if display_caption and CATEGORIES:
txt = '#{} {}'.format(int(bbox['track_id']),
CATEGORIES[int(bbox['category_id'])])
t_size = cv2.getTextSize(txt, cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
cv2.putText(image, txt,
(int(bbox['bbox']['x'] + bbox['bbox']['w'] / 2 -
t_size[0] / 2),
int(bbox['bbox']['y'] + t_size[1] + 10)),
cv2.FONT_HERSHEY_PLAIN, 1.5, (255, 255, 255), 2)
if dir_output is not None:
cv2.imwrite(dir_output, image)
return image
def visualize(dir_root, display=False, category_file=None, videonames=[]):
info = load_json()
palette = create_palette()
colors = {}
for x in palette:
colors[x['id']] = x['rgb']
if category_file is not None and os.path.isfile(category_file):
data = load_json(category_file)
data = data['category']
categories = ['__background']
categories.extend([cat['name'] for cat in data])
else:
categories = CATEGORIES
if not videonames:
videonames = os.listdir(os.path.join(info['dataset_dir'], 'Images'))
videonames = sorted(videonames)
tq = tqdm.tqdm(total=len(videonames))
alpha = 0.6
for videoname in videonames:
tq.set_description('Video {}'.format(videoname))
tq.update(1)
dir_input = os.path.join(dir_root, videoname)
dir_output = os.path.join(dir_root, '..', 'Visualization', videoname)
make_dir(dir_output)
files = os.listdir(
os.path.join(info['dataset_dir'], 'Images', videoname))
files = sorted(files)
for file in files:
filename = os.path.splitext(file)[0]
ori_im = cv2.imread(
os.path.join(info['dataset_dir'], 'Images', videoname, file))
image = ori_im.copy()
if os.path.isfile(os.path.join(dir_input, filename + '.json')):
with open(os.path.join(dir_input, filename + '.json')) as f:
data = json.load(f)
if isinstance(data, dict):
data = data.values()
for bbox in data:
top_left = (int(bbox['bbox']['x']), int(bbox['bbox']['y']))
bottom_right = (int(bbox['bbox']['x']) +
int(bbox['bbox']['w']),
int(bbox['bbox']['y']) +
int(bbox['bbox']['h']))
image = cv2.rectangle(image, tuple(top_left),
tuple(bottom_right),
tuple(colors[bbox['category_id']]),
1)
label = categories[int(bbox['category_id'])]
txt = '{}:{}:{}'.format(label, bbox['track_id'],
round(bbox['score'] * 100))
t_size = cv2.getTextSize(txt, cv2.FONT_HERSHEY_PLAIN, 1,
2)[0]
image = cv2.rectangle(
image, tuple(top_left),
tuple((top_left[0] + t_size[0] + 5,
top_left[1] + int(t_size[1] * 1) + 10)),
tuple(colors[bbox['category_id']]), -1)
cv2.putText(
image, txt,
(int(top_left[0] + 5), top_left[1] + t_size[1] + 5),
cv2.FONT_HERSHEY_PLAIN, 1, [0, 0, 0], 2)
cv2.imwrite(os.path.join(dir_output, filename + '.jpg'), image)
if display:
image = cv2.resize(
image, (int(image.shape[1] / 2), int(image.shape[0] / 2)))
cv2.imshow('demo', image)
cv2.waitKey(0)