Spaces:
Running
Running
File size: 3,355 Bytes
a383d0e |
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 |
from tqdm import tqdm
import json
import cv2
from os.path import join as pjoin
from config.CONFIG_UIED import Config
C = Config()
def draw_bounding_box_class(org, components, color=C.COLOR, line=2, show=False, write_path=None):
"""
Draw bounding box of components with their classes on the original image
:param org: original image
:param components: bbox [(column_min, row_min, column_max, row_max)]
-> top_left: (column_min, row_min)
-> bottom_right: (column_max, row_max)
:param color_map: colors mapping to different components
:param line: line thickness
:param compo_class: classes matching the corners of components
:param show: show or not
:return: labeled image
"""
board = org.copy()
bboxes = components['bboxes']
categories = components['categories']
for i in range(len(bboxes)):
bbox = bboxes[i]
category = categories[i]
board = cv2.rectangle(board, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color[C.CLASS_MAP[str(category)]], line)
board = cv2.putText(board, C.CLASS_MAP[str(category)], (bbox[0]+5, bbox[1]+20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color[C.CLASS_MAP[str(category)]], 2)
if show:
cv2.imshow('a', cv2.resize(board, (500, 1000)))
cv2.waitKey(0)
if write_path is not None:
cv2.imwrite(write_path, board)
return board
def load_ground_truth_json(gt_file, no_text=True):
def get_img_by_id(img_id):
for image in images:
if image['id'] == img_id:
return image['file_name'].split('/')[-1][:-4], (image['height'], image['width'])
def cvt_bbox(bbox):
'''
:param bbox: [x,y,width,height]
:return: [col_min, row_min, col_max, row_max]
'''
bbox = [int(b) for b in bbox]
return [bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]]
data = json.load(open(gt_file, 'r'))
images = data['images']
annots = data['annotations']
compos = {}
print('Loading %d ground truth' % len(annots))
for annot in tqdm(annots):
img_name, size = get_img_by_id(annot['image_id'])
if no_text and int(annot['category_id']) == 14:
compos[img_name] = {'bboxes': [], 'categories': [], 'size': size}
continue
if img_name not in compos:
compos[img_name] = {'bboxes': [cvt_bbox(annot['bbox'])], 'categories': [annot['category_id']], 'size':size}
else:
compos[img_name]['bboxes'].append(cvt_bbox(annot['bbox']))
compos[img_name]['categories'].append(annot['category_id'])
return compos
def view_gt_all(gt, img_root):
for img_id in gt:
compos = gt[img_id]
img = cv2.imread(pjoin(img_root, img_id + '.jpg'))
print(pjoin(img_root, img_id + '.jpg'))
draw_bounding_box_class(img, compos, show=True)
def view_gt_single(gt, img_root, img_id):
img_id = str(img_id)
compos = gt[img_id]
img = cv2.imread(pjoin(img_root, img_id + '.jpg'))
print(pjoin(img_root, img_id + '.jpg'))
draw_bounding_box_class(img, compos, show=True)
gt = load_ground_truth_json('E:\\Mulong\\Datasets\\rico\\instances_test.json', no_text=False)
# view_gt_all(gt, 'E:\\Mulong\\Datasets\\rico\\combined')
view_gt_single(gt, 'E:\\Mulong\\Datasets\\rico\\combined', 670)
|