Spaces:
Runtime error
Runtime error
| """ | |
| Utilities for polygon manipulation. | |
| """ | |
| import numpy as np | |
| import torch | |
| def is_clockwise(points): | |
| """Check whether a sequence of points is clockwise ordered""" | |
| # points is a list of 2d points. | |
| assert len(points) > 0 | |
| s = 0.0 | |
| for p1, p2 in zip(points, points[1:] + [points[0]]): | |
| s += (p2[0] - p1[0]) * (p2[1] + p1[1]) | |
| return s > 0.0 | |
| def resort_corners(corners): | |
| """Resort a sequence of corners so that the first corner starts | |
| from upper-left and counterclockwise ordered in image | |
| """ | |
| corners = corners.reshape(-1, 2) | |
| x_y_square_sum = corners[:, 0] ** 2 + corners[:, 1] ** 2 | |
| start_corner_idx = np.argmin(x_y_square_sum) | |
| corners_sorted = np.concatenate([corners[start_corner_idx:], corners[:start_corner_idx]]) | |
| ## sort points clockwise (counterclockwise in image) | |
| if not is_clockwise(corners_sorted[:, :2].tolist()): | |
| corners_sorted[1:] = np.flip(corners_sorted[1:], 0) | |
| return corners_sorted.reshape(-1) | |
| def get_all_order_corners(corners): | |
| """Get all possible permutation of a polygon""" | |
| length = int(len(corners) / 2) | |
| all_corners = torch.stack([corners.roll(i * 2) for i in range(length)]) | |
| return all_corners | |
| def pad_gt_polys(gt_instances, num_queries_per_poly, image_size, drop_rate, device=None): | |
| """Pad the ground truth polygons so that they have a uniform length""" | |
| thr_length = num_queries_per_poly * 2 | |
| room_targets = [] | |
| # padding ground truth on-fly | |
| for gt_inst in gt_instances: | |
| room_dict = {} | |
| room_corners = [] | |
| corner_labels = [] | |
| corner_lengths = [] | |
| for i, poly in enumerate(gt_inst.gt_masks.polygons): | |
| corners = torch.from_numpy(poly[0]).to(device) | |
| corners = torch.clip(corners, 0, image_size - 1) / (image_size - 1) | |
| # automatically skip the polygon if it is too long | |
| if len(corners) > thr_length: | |
| continue | |
| corner_lengths.append(len(corners)) | |
| corners_pad = torch.zeros(num_queries_per_poly * 2, device=device) | |
| corners_pad[: len(corners)] = corners | |
| labels = torch.ones(int(len(corners) / 2), dtype=torch.int64).to(device) | |
| labels_pad = torch.zeros(num_queries_per_poly, device=device) | |
| labels_pad[: len(labels)] = labels | |
| room_corners.append(corners_pad) | |
| corner_labels.append(labels_pad) | |
| room_classes = gt_inst.gt_classes | |
| if drop_rate > 0.0: | |
| keep_indices = np.where(np.random.rand(len(room_corners)) >= drop_rate)[0].tolist() | |
| if len(keep_indices) > 0: # Only apply drop if we have something left | |
| room_corners = [room_corners[i] for i in keep_indices] | |
| corner_labels = [corner_labels[i] for i in keep_indices] | |
| corner_lengths = [corner_lengths[i] for i in keep_indices] | |
| room_classes = gt_inst.gt_classes[keep_indices] | |
| room_dict = { | |
| "coords": torch.stack(room_corners).to(device), | |
| "labels": torch.stack(corner_labels).to(device), | |
| "lengths": torch.tensor(corner_lengths, device=device), | |
| "room_labels": room_classes.to(device), | |
| } | |
| room_targets.append(room_dict) | |
| return room_targets | |