Spaces:
Runtime error
Runtime error
| # ------------------------------------------------------------------------ | |
| # Copyright (c) 2023-present, BAAI. All Rights Reserved. | |
| # | |
| # 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. | |
| # ------------------------------------------------------------------------ | |
| """Prompt encoder.""" | |
| import torch | |
| from torch import nn | |
| class PromptEncoder(nn.Module): | |
| """Module to encode geometric prompts.""" | |
| def __init__(self, embed_dim, image_size): | |
| super(PromptEncoder, self).__init__() | |
| self.point_embed = nn.Embedding(5, embed_dim) # [bg, fg, lt, rb, pad] | |
| self.corner_labels = torch.tensor([[2, 3]], dtype=torch.int64) | |
| self.register_buffer("coord_matrix", torch.randn((2, embed_dim // 2))) | |
| self.img_pos, self.img_size = None, [image_size] * 2 | |
| def as_tensor(self, input): | |
| """Convert input into a tensor.""" | |
| return torch.as_tensor(input, device=self.coord_matrix.device) | |
| def to_points(self, points=None, boxes=None): | |
| """Convert points or boxes to point prompts.""" | |
| if points is not None: | |
| if isinstance(points, (tuple, list)): | |
| coords, labels = points | |
| else: | |
| coords, labels = points[:, :, :2], points[:, :, 2] | |
| coords = coords.__add__(0.5).__itruediv__(self.img_size[::-1]) | |
| coords = self.as_tensor(coords.clip(0, 1).astype("float32")) | |
| labels = self.as_tensor(labels.astype("int64")) | |
| return coords, labels | |
| if boxes is not None: | |
| coords = boxes.reshape((-1, 2, 2)) | |
| coords = coords.__add__(0.5).__itruediv__(self.img_size[::-1]) | |
| coords = self.as_tensor(coords.clip(0, 1).astype("float32")) | |
| labels = self.as_tensor(self.corner_labels) | |
| return coords, labels | |
| return None | |
| def encode_coords(self, coords): | |
| """Return the embedding for given coords.""" | |
| pi4, pi2 = 4 * 3.1415926, 2 * 3.1415926 | |
| if self.coord_matrix.dtype != torch.float32: | |
| self.coord_matrix = self.coord_matrix.float() | |
| rad = coords.mul(pi4).sub_(pi2) @ self.coord_matrix | |
| dtype = self.point_embed.weight.dtype | |
| return torch.cat([rad.sin(), rad.cos()], dim=-1).to(dtype=dtype) | |
| def encode_points(self, coords, labels): | |
| """Return the embedding for given points.""" | |
| embed = self.encode_coords(coords) | |
| embed.mul_(labels.ne(4).unsqueeze_(-1).float().to(dtype=embed.dtype)) | |
| return embed.add_(self.point_embed(labels)) | |
| def encode_grid(self, grid_size): | |
| """Return the embedding for a grid of specified size.""" | |
| grid = torch.ones(*grid_size, dtype=torch.float32) | |
| y = grid.cumsum(dim=0).sub_(0.5).div_(grid_size[0]) | |
| x = grid.cumsum(dim=1).sub_(0.5).div_(grid_size[1]) | |
| coords = self.as_tensor(torch.stack([x, y], dim=-1)) | |
| return self.encode_coords(coords) | |
| def forward(self, inputs): | |
| sparse_embeds = [] | |
| if inputs.get("boxes", None) is not None: | |
| coords, labels = self.to_points(boxes=inputs["boxes"]) | |
| sparse_embeds.append(self.encode_points(coords, labels)) | |
| if inputs.get("points", None) is not None: | |
| coords, labels = self.to_points(points=inputs["points"]) | |
| sparse_embeds.append(self.encode_points(coords, labels)) | |
| if len(sparse_embeds) > 1: | |
| sparse_embeds = [torch.cat(sparse_embeds, dim=1)] | |
| elif len(sparse_embeds) == 0: | |
| raise ValueError("Excepted ``points`` or ``boxes`` prompts.") | |
| img_embed_size = torch.Size(inputs["img_embeds"].shape[2:-1]) | |
| if self.img_pos is None or self.img_pos.shape[0] != img_embed_size.numel(): | |
| self.img_pos = self.encode_grid(img_embed_size).flatten(0, 1) | |
| return {"sparse_embeds": sparse_embeds[0], "img_pos": self.img_pos} | |