Spaces:
Sleeping
Sleeping
| import re | |
| import numpy as np | |
| from PIL import ImageDraw | |
| def parse_bbox_and_labels(detokenized_output: str): | |
| matches = re.finditer( | |
| '<loc(?P<y0>\d\d\d\d)><loc(?P<x0>\d\d\d\d)><loc(?P<y1>\d\d\d\d)><loc(?P<x1>\d\d\d\d)>' | |
| ' (?P<label>.+?)( ;|$)', | |
| detokenized_output, | |
| ) | |
| labels, boxes = [], [] | |
| fmt = lambda x: float(x) / 1024.0 | |
| for m in matches: | |
| d = m.groupdict() | |
| boxes.append([fmt(d['y0']), fmt(d['x0']), fmt(d['y1']), fmt(d['x1'])]) | |
| labels.append(d['label']) | |
| return np.array(boxes), np.array(labels) | |
| def display_boxes(image, boxes, labels, target_size): | |
| h, w = target_size | |
| # fig, ax = plt.subplots() | |
| # ax.imshow(image) | |
| draw = ImageDraw.Draw(image) | |
| for i in range(boxes.shape[0]): | |
| y, x, y2, x2 = (boxes[i][0]*w,boxes[i][1]*h,boxes[i][2]*w,boxes[i][3]*h) | |
| # width = x2 - x | |
| # height = y2 - y | |
| # Create a Rectangle patch | |
| # rect = patches.Rectangle((x, y), | |
| # width, | |
| # height, | |
| # linewidth=1, | |
| # edgecolor='r', | |
| # facecolor='none') | |
| draw.rectangle((x,y,x2,y2) , outline="red", width=3) | |
| # Add label | |
| # plt.text(x, y, labels[i], color='red', fontsize=12) | |
| # # Add the patch to the Axes | |
| # ax.add_patch(rect) | |
| # plt.show() | |
| return image |