Spaces:
Runtime error
Runtime error
| import numpy as np | |
| def expand_bbox_to_ratio(bbox, imshape, target_aspect_ratio): | |
| x0, y0, x1, y1 = [int(_) for _ in bbox] | |
| h, w = y1 - y0, x1 - x0 | |
| cur_ratio = h / w | |
| if cur_ratio == target_aspect_ratio: | |
| return [x0, y0, x1, y1] | |
| if cur_ratio < target_aspect_ratio: | |
| target_height = int(w*target_aspect_ratio) | |
| y0, y1 = expand_axis(y0, y1, target_height, imshape[0]) | |
| else: | |
| target_width = int(h/target_aspect_ratio) | |
| x0, x1 = expand_axis(x0, x1, target_width, imshape[1]) | |
| return x0, y0, x1, y1 | |
| def expand_axis(start, end, target_width, limit): | |
| # Can return a bbox outside of limit | |
| cur_width = end - start | |
| start = start - (target_width-cur_width)//2 | |
| end = end + (target_width-cur_width)//2 | |
| if end - start != target_width: | |
| end += 1 | |
| assert end - start == target_width | |
| if start < 0 and end > limit: | |
| return start, end | |
| if start < 0 and end < limit: | |
| to_shift = min(0 - start, limit - end) | |
| start += to_shift | |
| end += to_shift | |
| if end > limit and start > 0: | |
| to_shift = min(end - limit, start) | |
| end -= to_shift | |
| start -= to_shift | |
| assert end - start == target_width | |
| return start, end | |
| def expand_box(bbox, imshape, mask, percentage_background: float): | |
| assert isinstance(bbox[0], int) | |
| assert 0 < percentage_background < 1 | |
| # Percentage in S | |
| mask_pixels = mask.long().sum().cpu() | |
| total_pixels = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) | |
| percentage_mask = mask_pixels / total_pixels | |
| if (1 - percentage_mask) > percentage_background: | |
| return bbox | |
| target_pixels = mask_pixels / (1 - percentage_background) | |
| x0, y0, x1, y1 = bbox | |
| H = y1 - y0 | |
| W = x1 - x0 | |
| p = np.sqrt(target_pixels/(H*W)) | |
| target_width = int(np.ceil(p * W)) | |
| target_height = int(np.ceil(p * H)) | |
| x0, x1 = expand_axis(x0, x1, target_width, imshape[1]) | |
| y0, y1 = expand_axis(y0, y1, target_height, imshape[0]) | |
| return [x0, y0, x1, y1] | |
| def expand_axises_by_percentage(bbox_XYXY, imshape, percentage): | |
| x0, y0, x1, y1 = bbox_XYXY | |
| H = y1 - y0 | |
| W = x1 - x0 | |
| expansion = int(((H*W)**0.5) * percentage) | |
| new_width = W + expansion | |
| new_height = H + expansion | |
| x0, x1 = expand_axis(x0, x1, min(new_width, imshape[1]), imshape[1]) | |
| y0, y1 = expand_axis(y0, y1, min(new_height, imshape[0]), imshape[0]) | |
| return [x0, y0, x1, y1] | |
| def get_expanded_bbox( | |
| bbox_XYXY, | |
| imshape, | |
| mask, | |
| percentage_background: float, | |
| axis_minimum_expansion: float, | |
| target_aspect_ratio: float): | |
| bbox_XYXY = bbox_XYXY.long().cpu().numpy().tolist() | |
| # Expand each axis of the bounding box by a minimum percentage | |
| bbox_XYXY = expand_axises_by_percentage(bbox_XYXY, imshape, axis_minimum_expansion) | |
| # Find the minimum bbox with the aspect ratio. Can be outside of imshape | |
| bbox_XYXY = expand_bbox_to_ratio(bbox_XYXY, imshape, target_aspect_ratio) | |
| # Expands square box such that X% of the bbox is background | |
| bbox_XYXY = expand_box(bbox_XYXY, imshape, mask, percentage_background) | |
| assert isinstance(bbox_XYXY[0], (int, np.int64)) | |
| return bbox_XYXY | |
| def include_box(bbox, minimum_area, aspect_ratio_range, min_bbox_ratio_inside, imshape): | |
| def area_inside_ratio(bbox, imshape): | |
| area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) | |
| area_inside = (min(bbox[2], imshape[1]) - max(0, bbox[0])) * (min(imshape[0], bbox[3]) - max(0, bbox[1])) | |
| return area_inside / area | |
| ratio = (bbox[3] - bbox[1]) / (bbox[2] - bbox[0]) | |
| area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0]) | |
| if area_inside_ratio(bbox, imshape) < min_bbox_ratio_inside: | |
| return False | |
| if ratio <= aspect_ratio_range[0] or ratio >= aspect_ratio_range[1] or area < minimum_area: | |
| return False | |
| return True | |