Spaces:
Runtime error
Runtime error
| # coding=utf-8 | |
| # Copyright 2024 The Google Research Authors. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # Copyright 2024 The Google Research Authors. | |
| # This file is based on the SAM (Segment Anything) and HQ-SAM. | |
| # | |
| # https://github.com/facebookresearch/segment-anything | |
| # https://github.com/SysCV/sam-hq/tree/main | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """SAM Utilities.""" | |
| # pylint: disable=all | |
| # pylint: disable=g-importing-member | |
| import json | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| from scipy.spatial.distance import cdist | |
| def show_mask(mask, ax, random_color=False): | |
| if random_color: | |
| color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0) | |
| else: | |
| color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6]) | |
| h, w = mask.shape[-2:] | |
| mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) | |
| ax.imshow(mask_image) | |
| def show_points(coords, labels, ax, marker_size=375): | |
| pos_points = coords[labels == 1] | |
| neg_points = coords[labels == 0] | |
| ax.scatter( | |
| pos_points[:, 0], | |
| pos_points[:, 1], | |
| color='green', | |
| marker='*', | |
| s=marker_size, | |
| edgecolor='white', | |
| linewidth=1.25, | |
| ) | |
| ax.scatter( | |
| neg_points[:, 0], | |
| neg_points[:, 1], | |
| color='red', | |
| marker='*', | |
| s=marker_size, | |
| edgecolor='white', | |
| linewidth=1.25, | |
| ) | |
| def show_box(box, ax): | |
| x0, y0, x1, y1 = box | |
| w, h = x1 - x0, y1 - y0 | |
| ax.add_patch( | |
| plt.Rectangle( | |
| (x0, y0), w, h, edgecolor='red', facecolor=(0, 0, 0, 0), lw=2 | |
| ) | |
| ) | |
| def show_anns(anns): | |
| if len(anns) == 0: | |
| return | |
| for index, dictionary in enumerate(anns): | |
| dictionary['id'] = index | |
| sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True) | |
| ax = plt.gca() | |
| ax.set_autoscale_on(False) | |
| # polygons = [] | |
| # color = [] | |
| for ann in sorted_anns: | |
| m = ann['segmentation'] | |
| img = np.ones((m.shape[0], m.shape[1], 3)) | |
| color_mask = np.random.random((1, 3)).tolist()[0] | |
| for i in range(3): | |
| img[:, :, i] = color_mask[i] | |
| ax.imshow(np.dstack((img, m * 0.35))) | |
| # Get the centroid of the mask | |
| mask_y, mask_x = np.nonzero(m) | |
| centroid_x, centroid_y = np.mean(mask_x), np.mean(mask_y) | |
| # Display the mask ID | |
| mask_id = ann['id'] | |
| ax.text( | |
| centroid_x, | |
| centroid_y, | |
| str(mask_id), | |
| color='black', | |
| fontsize=48, | |
| weight='bold', | |
| ) | |
| # Turn CAM result to SAM prompt | |
| def aggregate_RGB_channel(activation_mask, how='max'): | |
| B, C, H, W = activation_mask.shape | |
| if how == 'max': | |
| res_activation_mask = np.amax(activation_mask, axis=1, keepdims=True) | |
| elif how == 'avr': | |
| res_activation_mask = np.mean(activation_mask, axis=1, keepdims=True) | |
| res_activation_mask = res_activation_mask.reshape(B, 1, H * W) | |
| res_activation_mask = np.squeeze(res_activation_mask, axis=1) | |
| return res_activation_mask | |
| def find_k_points(arr, k, order='max', how_filter='median'): | |
| arr = arr.squeeze(0) | |
| flat_indices = np.argpartition(arr.flatten(), -k)[-k:] | |
| unravel_topk_idx = np.unravel_index(flat_indices, arr.shape) | |
| topk_indices = np.array(unravel_topk_idx).transpose()[:, ::-1] | |
| # print(topk_indices.shape) | |
| if how_filter == 'random': | |
| random_rows = np.random.choice( | |
| topk_indices.shape[0], size=int(round(k / 16)), replace=False | |
| ) | |
| topk_indices = topk_indices[random_rows] | |
| elif how_filter == 'median': | |
| distances = cdist(topk_indices, topk_indices) | |
| distances = np.sum(distances, axis=1) | |
| median_distance = np.median(distances) | |
| filtered_idx = [ | |
| i for i in range(len(distances)) if distances[i] < median_distance | |
| ] | |
| topk_indices = topk_indices[filtered_idx] | |
| return topk_indices | |
| def max_sum_submatrix(matrix): | |
| matrix = np.array(matrix) | |
| H, W = matrix.shape | |
| # Preprocess cumulative sums for rows | |
| matrix[:, 1:] += matrix[:, :-1] | |
| max_sum = float('-inf') | |
| max_rect = (0, 0, 0, 0) # (top, left, bottom, right) | |
| for left in range(W): | |
| for right in range(left, W): | |
| # Apply 1D Kadane's algorithm for the current pair of columns | |
| column_sum = matrix[:, right] - (matrix[:, left - 1] if left > 0 else 0) | |
| max_ending_here = max_so_far = column_sum[0] | |
| start, end = 0, 0 | |
| for i in range(1, H): | |
| val = column_sum[i] | |
| if max_ending_here > 0: | |
| max_ending_here += val | |
| else: | |
| max_ending_here = val | |
| start = i | |
| if max_ending_here > max_so_far: | |
| max_so_far = max_ending_here | |
| end = i | |
| if max_so_far > max_sum: | |
| max_sum = max_so_far | |
| max_rect = (start, left, end, right) | |
| return max_sum, max_rect | |
| def CAM2SAMClick(activation_map, k=5, order='max', how_filter='median'): | |
| # activation_map = aggregate_RGB_channel(activation_map) | |
| H, W, C = activation_map.shape | |
| activation_map = activation_map.reshape((1, 1, H, W)) | |
| coords = [] | |
| for nrow in range(activation_map.shape[0]): | |
| coord = find_k_points(activation_map[nrow], k, order, how_filter) | |
| coords.append(coord) | |
| return coords | |
| def CAM2SAMBox(activation_map): | |
| # print(activation_map.shape) | |
| # activation_map = aggregate_RGB_channel(activation_map) | |
| H, W, C = activation_map.shape | |
| activation_map = activation_map.reshape((1, H, W)) | |
| box_coordinates = [] | |
| for nrow in range(activation_map.shape[0]): | |
| # print(activation_map[nrow].shape) | |
| arr = activation_map[nrow] | |
| norm_arr = 2 * ((arr - np.min(arr)) / (np.max(arr) - np.min(arr))) - 1 | |
| # print(norm_arr.shape) | |
| _, box_coordinate = max_sum_submatrix(norm_arr) | |
| box_coordinates.append(box_coordinate) | |
| return box_coordinates | |
| # Visualize | |
| def visualize_attention(arr, filename): | |
| # Create a figure and axes object | |
| fig, ax = plt.subplots() | |
| # Display the array as an image | |
| im = ax.imshow(arr) | |
| # Add a colorbar | |
| ax.figure.colorbar(im, ax=ax) | |
| # cbar = ax.figure.colorbar(im, ax=ax) | |
| # Save the figure as a PNG file | |
| fig.savefig(filename) | |
| # Build config | |
| # def build_sam_config(config_path): | |
| # with open(config_path, 'r') as infile: | |
| # config = json.load(infile) | |
| # sam_checkpoint = config['model']['sam_checkpoint'] | |
| # model_type = config['model']['model_type'] | |
| # return sam_checkpoint, model_type | |
| def build_sam_config(config): | |
| sam_checkpoint = config.sam.sam_checkpoint | |
| model_type = config.sam.model_type | |
| return sam_checkpoint, model_type |