import os import math import logging import numpy as np import kornia import torch import torch.nn as nn import torch.nn.functional as F try: from .task_utils import CenterPadding, upsample_features except ImportError: from task_utils import CenterPadding, upsample_features logging.getLogger().setLevel(logging.WARNING) class FeatureExtractor(nn.Module): def __init__(self, config, device, return_class_token=False): super(FeatureExtractor, self).__init__() self.feature_extractor = config['pretrained']['feature_extractor'] self.patch_size = config['architecture']['patch_size'] self.backbone_ckpt_path = os.path.join(config['logging']['save_dir'], config['logging']['exp_name'], 'dinov3_vitl16_pretrain_lvd1689m-8aa4cbdd.pth') self.head_ckpt_path = os.path.join(config['logging']['save_dir'], config['logging']['exp_name'], 'dinov3_vitl16_dinotxt_vision_head_and_text_encoder-a442d8f5.pth') self.return_class_token = return_class_token self.device = device if self.feature_extractor == 'dinov3_vitl16': self.backbone, _ = torch.hub.load('facebookresearch/dinov3', 'dinov3_vitl16_dinotxt_tet1280d20h24l', backbone_weights=self.backbone_ckpt_path, weights=self.head_ckpt_path) del self.backbone.text_model self.model = self.backbone.visual_model.backbone.to(device) self.head = self.backbone.visual_model.head.to(device) else: raise ValueError(f'Feature extractor {self.feature_extractor} not supported.') def extract_dinov3(self, images, batch_size=1024, patch_length=16, layers=[23]): transform = kornia.augmentation.AugmentationSequential( CenterPadding(multiple=patch_length), kornia.augmentation.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ) transformed_images = transform(images) class_tokens, text_aligned_class_tokens, patch_tokens, register_tokens, feature_maps = [], [], [], [], [] for i in range(0, transformed_images.shape[0], batch_size): image_batch = transformed_images[i:(i + batch_size)].to(device=self.device) with torch.inference_mode(): features_out = self.model.get_intermediate_layers(image_batch, return_class_token=True, return_extra_tokens=True, n=layers) class_tokens.append(features_out[-1][1]) patch_tokens.append(features_out[0][0]) register_tokens.append(features_out[0][2]) head_inputs = torch.cat([class_tokens[-1].unsqueeze(1), register_tokens[-1], patch_tokens[-1]], dim=1) text_aligned_class_token = self.head(head_inputs)[0][0:1] text_aligned_class_tokens.append(text_aligned_class_token) B, _, C = patch_tokens[-1].size() H, W = image_batch.shape[2], image_batch.shape[3] patch_H, patch_W = math.ceil(H / patch_length), math.ceil(W / patch_length) feature_maps.append(patch_tokens[-1].permute(0, 2, 1).view(B, C, patch_H, patch_W)) class_tokens = torch.cat(class_tokens, dim=0) text_aligned_class_tokens = torch.cat(text_aligned_class_tokens, dim=0) patch_tokens = torch.cat(patch_tokens, dim=0) register_tokens = torch.cat(register_tokens, dim=0) feature_maps = torch.cat(feature_maps, dim=0) return class_tokens, text_aligned_class_tokens, patch_tokens, register_tokens, feature_maps def forward(self, images, resize=False): if self.feature_extractor == 'dinov3_vitl16': class_tokens, text_aligned_class_tokens, patch_tokens, register_tokens, feature_maps = self.extract_dinov3(images) if resize: image_height, image_width = images.shape[2], images.shape[3] padded_height = math.ceil(image_height / self.patch_size) * self.patch_size padded_width = math.ceil(image_width / self.patch_size) * self.patch_size resized_feature_maps = [] chunk_size = 32 for i in range(0, len(feature_maps), chunk_size): resized_feature_maps.append(upsample_features(feature_maps[i:i + chunk_size], image_height, image_width, padded_height, padded_width)) feature_maps = torch.cat(resized_feature_maps) return { 'class_tokens': class_tokens, 'text_aligned_class_tokens': text_aligned_class_tokens, 'patch_tokens': patch_tokens, 'register_tokens': register_tokens, 'feature_maps': feature_maps } class RegionTokenGenerator(nn.Module): def __init__(self, pooling_method='average', device='cuda'): super(RegionTokenGenerator, self).__init__() self.model, _ = torch.hub.load('facebookresearch/dinov3', 'dinov3_vitl16_dinotxt_tet1280d20h24l') self.model = self.model.visual_model.head.to(device) self.pooling_method = pooling_method self.device = device def forward(self, regions, class_tokens, feature_maps, register_tokens): pooled_tokens, text_aligned_tokens = [], [] for scale_idx in range(regions.shape[2]): scale_pooled_tokens, scale_text_aligned_tokens = [], [] for batch_idx in range(len(regions)): image_regions = regions[batch_idx, :, scale_idx] image_class_tokens = class_tokens[batch_idx] image_feature_maps = feature_maps[batch_idx] image_register_tokens = register_tokens[batch_idx] if image_regions.numel() == 0: scale_pooled_tokens.append(torch.zeros((0, image_feature_maps.shape[0]), device=self.device)) scale_text_aligned_tokens.append(torch.zeros((0, image_feature_maps.shape[0]), device=self.device)) continue # Get the features that pertain to the regions region_features = torch.einsum('rhw,chw->rc', image_regions.float(), image_feature_maps) text_alignment_inputs = torch.cat([image_class_tokens[None], image_register_tokens, region_features], dim=0) text_aligned_region_features = self.model(text_alignment_inputs[None])[0][image_register_tokens.shape[0] + 1 :] # Pool the region features if self.pooling_method == 'average': valid_elements = image_regions.sum(dim=(1, 2), dtype=torch.float32).clamp(min=1).unsqueeze(1) region_features = region_features / valid_elements text_aligned_region_features = text_aligned_region_features / valid_elements else: raise ValueError(f'Pooling method {self.pooling_method} not supported.') scale_pooled_tokens.append(region_features) scale_text_aligned_tokens.append(text_aligned_region_features) scale_pooled_tokens = torch.stack(scale_pooled_tokens) scale_text_aligned_tokens = torch.stack(scale_text_aligned_tokens) pooled_tokens.append(scale_pooled_tokens.unsqueeze(2)) text_aligned_tokens.append(scale_text_aligned_tokens.unsqueeze(2)) pooled_tokens = torch.cat(pooled_tokens, dim=2) text_aligned_tokens = torch.cat(text_aligned_tokens, dim=2) return { 'pooled_tokens': pooled_tokens, 'text_aligned_tokens': text_aligned_tokens } class TextEncoder(nn.Module): def __init__(self, config, device, prompt='photo of a '): super(TextEncoder, self).__init__() self.ckpt_path = os.path.join(config['logging']['save_dir'], config['logging']['exp_name'], 'dinov3_vitl16_dinotxt_vision_head_and_text_encoder-a442d8f5.pth') self.model, self.tokenizer = torch.hub.load('facebookresearch/dinov3', 'dinov3_vitl16_dinotxt_tet1280d20h24l', weights=self.ckpt_path) del self.model.visual_model self.model = self.model.to(device) self.prompt = prompt self.device = device def forward(self, texts): texts = [self.prompt + t for t in texts] text_tokens = self.tokenizer.tokenize(texts).to(self.device) with torch.no_grad(): text_embeddings = self.model.encode_text(text_tokens) text_embeddings = text_embeddings[:, text_embeddings.shape[1] // 2 :] return text_embeddings class PositionalEmbedding2D(nn.Module): def __init__(self, embedding_dim=64, scale=None): super().__init__() if scale is None or scale <= 0.0: scale = 1.0 generator = torch.Generator() generator.manual_seed(42) self.register_buffer("positional_encoding_gaussian_matrix", scale * torch.randn((2, embedding_dim // 2), generator=generator)) def _pe_encoding(self, coords): coords = 2 * coords - 1 coords = coords @ self.positional_encoding_gaussian_matrix coords = 2 * np.pi * coords return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) def forward(self, size): h, w = size device = self.positional_encoding_gaussian_matrix.device grid = torch.ones((h, w), device=device, dtype=torch.float32) y_embed = grid.cumsum(dim=0) - 0.5 x_embed = grid.cumsum(dim=1) - 0.5 y_embed = y_embed / h x_embed = x_embed / w pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) return pe.permute(2, 0, 1) class AttentionLayer(nn.Module): def __init__(self, q_dim, kv_dim, hidden_dim, num_heads=8, dropout=0.1, use_bias=False, use_v_proj=True, use_out_proj=True): super(AttentionLayer, self).__init__() self.hidden_dim = hidden_dim self.num_heads = num_heads assert hidden_dim % num_heads == 0, 'Hidden dimension must be a multiple of the number of heads.' self.head_dim = hidden_dim // num_heads if not use_v_proj: assert kv_dim == hidden_dim, 'Key and value dimensions must be the same as the hidden dimension if not using v_proj.' self.q_proj = nn.Linear(q_dim, hidden_dim, bias=use_bias) nn.init.kaiming_normal_(self.q_proj.weight, mode='fan_in', nonlinearity='linear') self.k_proj = nn.Linear(kv_dim, hidden_dim, bias=use_bias) nn.init.kaiming_normal_(self.k_proj.weight, mode='fan_in', nonlinearity='linear') if use_v_proj: self.v_proj = nn.Linear(kv_dim, hidden_dim, bias=use_bias) nn.init.kaiming_normal_(self.v_proj.weight, mode='fan_in', nonlinearity='linear') else: self.v_proj = nn.Identity() if use_bias: nn.init.zeros_(self.q_proj.bias) nn.init.zeros_(self.k_proj.bias) if use_v_proj: nn.init.zeros_(self.v_proj.bias) self.q_norm = nn.LayerNorm(self.head_dim) self.k_norm = nn.LayerNorm(self.head_dim) self.dropout = nn.Dropout(dropout) if use_out_proj: self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=use_bias) nn.init.kaiming_normal_(self.out_proj.weight, mode='fan_in', nonlinearity='linear') if use_bias: nn.init.zeros_(self.out_proj.bias) else: self.out_proj = nn.Identity() self.scale = self.head_dim ** -0.5 def forward(self, q, k, v, mask=None, attn_threshold=None): batch_size, q_len, _ = q.shape _, kv_len, _ = k.shape query = self.q_proj(q).view(batch_size, q_len, self.num_heads, -1).transpose(1, 2) key = self.k_proj(k).view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2) value = self.v_proj(v).view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2) query = self.q_norm(query) key = self.k_norm(key) attn_scores = torch.matmul(query, key.transpose(-2, -1)) * self.scale if mask is not None: attn_scores = attn_scores.masked_fill(mask == 0, float('-inf')) if attn_threshold is not None: max_attn_scores, _ = attn_scores.max(dim=-1, keepdim=True) thresholding_mask = attn_scores >= (attn_threshold * max_attn_scores) attn_scores = attn_scores.masked_fill(thresholding_mask == 0, -1e9) attn_weights = F.softmax(attn_scores, dim=-1) attn_weights = self.dropout(attn_weights) attn_out = torch.matmul(attn_weights, value) attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, q_len, self.hidden_dim) out = self.out_proj(attn_out) return out, attn_weights class MLPBlock(nn.Module): def __init__(self, hidden_dim, intermediate_dim, dropout=0.1): super(MLPBlock, self).__init__() self.linear1 = nn.Linear(hidden_dim, intermediate_dim) self.gelu = nn.GELU() self.linear2 = nn.Linear(intermediate_dim, hidden_dim) self.dropout = nn.Dropout(dropout) def forward(self, x): z = self.linear1(x) z = self.gelu(z) z = self.dropout(z) z = self.linear2(z) return z class CrossAttentionBlock(nn.Module): def __init__(self, q_dim, kv_dim, hidden_dim, mlp_dim, num_heads, dropout, use_bias): super(CrossAttentionBlock, self).__init__() self.query_norm = nn.LayerNorm(q_dim) self.cross_attn = AttentionLayer(q_dim, kv_dim, hidden_dim, num_heads, dropout, use_bias) self.dropout = nn.Dropout(dropout) self.mlp_norm = nn.LayerNorm(hidden_dim) self.mlp = MLPBlock(hidden_dim, mlp_dim) self.out_norm = nn.LayerNorm(hidden_dim) def forward(self, query, context, mask=None): x = self.query_norm(query) x, attn_scores = self.cross_attn(q=x, k=context, v=context, mask=mask) x = self.dropout(x) x = x + query y = self.mlp_norm(x) y = self.mlp(y) out = self.out_norm(y) + x return out, attn_scores class TextAlignmentBlock(nn.Module): def __init__(self, hidden_dim, intermediate_dim, output_dim, dropout=0.1): super(TextAlignmentBlock, self).__init__() self.linear1 = nn.Linear(hidden_dim, intermediate_dim) self.gelu = nn.GELU() self.linear2 = nn.Linear(intermediate_dim, output_dim) self.dropout = nn.Dropout(dropout) def forward(self, x): z = self.linear1(x) z = self.gelu(z) z = self.dropout(z) z = self.linear2(z) return z class TokenAggregator(nn.Module): def __init__(self, config): super(TokenAggregator, self).__init__() self.merging_iou_threshold = config['parameters']['merging_iou_threshold'] self.merging_similarity_threshold = config['parameters']['merging_similarity_threshold'] self.binarization_threshold = config['parameters'].get('binarization_threshold', 0.5) def _compute_binary_masks(self, region_masks): num_masks = region_masks.shape[0] return (region_masks.reshape(num_masks, -1) > self.binarization_threshold).float() def _compute_iou_matrix(self, binary_masks): num_masks = binary_masks.shape[0] if num_masks == 0: return torch.zeros(0, 0, device=binary_masks.device) intersection = torch.mm(binary_masks, binary_masks.t()) areas = binary_masks.sum(dim=1) union = areas.unsqueeze(1) + areas.unsqueeze(0) - intersection return intersection / torch.clamp(union, min=1.0) def _find_connected_components(self, adjacency): n = adjacency.shape[0] if n == 0: return [] # Initialize labels labels = torch.arange(n, device=adjacency.device) # Iterative label propagation for _ in range(int(np.ceil(np.log2(n + 1))) + 1): neighbor_labels = torch.where(adjacency, labels.unsqueeze(0).expand(n, -1), labels.unsqueeze(1).expand(-1, n)) new_labels = neighbor_labels.min(dim=1)[0] new_labels = torch.minimum(new_labels, labels) if torch.equal(new_labels, labels): break labels = new_labels # Convert to groups labels_cpu = labels.cpu().tolist() label_to_group = {} for idx, label in enumerate(labels_cpu): if label not in label_to_group: label_to_group[label] = [] label_to_group[label].append(idx) return list(label_to_group.values()) def _compute_token_similarity_matrix(self, pred_tokens): num_tokens = pred_tokens.shape[0] if num_tokens == 0: return torch.zeros(0, 0, device=pred_tokens.device) pred_tokens = F.normalize(pred_tokens, p=2, dim=-1) return torch.mm(pred_tokens, pred_tokens.t()) def group_predictions(self, region_masks, pred_tokens=None): num_masks = region_masks.shape[0] if num_masks == 0: return [] binary_masks = self._compute_binary_masks(region_masks) iou_matrix = self._compute_iou_matrix(binary_masks) mask_adjacency = iou_matrix > self.merging_iou_threshold if pred_tokens is not None and pred_tokens.shape[0] == num_masks: token_sim = self._compute_token_similarity_matrix(pred_tokens) token_adjacency = token_sim > self.merging_similarity_threshold adjacency = mask_adjacency | token_adjacency else: adjacency = mask_adjacency return self._find_connected_components(adjacency) def forward(self, ren_outputs, remove_singleton_groups=True): pred_tokens = ren_outputs['pred_tokens'] region_masks = ren_outputs['region_masks'] text_aligned_tokens = ren_outputs['text_aligned_tokens'] pred_tokens = torch.flatten(pred_tokens, 1, 2) region_masks = torch.flatten(region_masks, 1, 2) text_aligned_tokens = torch.flatten(text_aligned_tokens, 1, 2) aggregated_outputs = {'pred_tokens': [], 'region_masks': [], 'text_aligned_tokens': []} for batch_idx in range(pred_tokens.shape[0]): batch_pred_tokens = pred_tokens[batch_idx] batch_region_masks = region_masks[batch_idx] batch_text_aligned_tokens = text_aligned_tokens[batch_idx] groups = self.group_predictions(batch_region_masks, batch_pred_tokens) kept_groups = [] for local_group_idxs in groups: if remove_singleton_groups and len(local_group_idxs) == 1: continue global_idxs = torch.tensor(local_group_idxs, device=batch_region_masks.device) group_mean_mask = batch_region_masks[global_idxs].mean(dim=0) kept_groups.append({'global_idxs': global_idxs, 'mean_mask': group_mean_mask}) if len(kept_groups) == 0: mask_areas = batch_region_masks.sum(dim=(-2, -1)) best_idx = mask_areas.argmax() kept_groups.append({ 'global_idxs': best_idx.unsqueeze(0), 'mean_mask': batch_region_masks[best_idx], }) new_pred_tokens, new_region_masks, new_text_aligned_tokens = [], [], [] for gd in kept_groups: global_idxs = gd['global_idxs'] new_pred_tokens.append(pred_tokens[batch_idx][global_idxs].mean(dim=0)) new_region_masks.append(gd['mean_mask']) new_text_aligned_tokens.append(text_aligned_tokens[batch_idx][global_idxs].mean(dim=0)) aggregated_outputs['pred_tokens'].append(torch.stack(new_pred_tokens, dim=0)) aggregated_outputs['region_masks'].append(torch.stack(new_region_masks, dim=0)) aggregated_outputs['text_aligned_tokens'].append(torch.stack(new_text_aligned_tokens, dim=0)) return aggregated_outputs class RegionEncoder(nn.Module): def __init__(self, config): super(RegionEncoder, self).__init__() hidden_dim = config['architecture']['hidden_dim'] text_embed_dim = config['architecture']['text_embed_dim'] image_resolution = config['parameters']['image_resolution'] patch_size = config['architecture']['patch_size'] feature_map_resolution = image_resolution // patch_size self.feature_map_resolution = feature_map_resolution self.image_resolution = image_resolution # Create position embeddings for the prompts and feature maps position_embedder = PositionalEmbedding2D(hidden_dim) location_embeddings = position_embedder((image_resolution, image_resolution)) feature_embeddings = position_embedder((feature_map_resolution, feature_map_resolution)).flatten(-2).permute(1, 0) self.register_buffer('location_embeddings', location_embeddings) self.register_buffer('feature_embeddings', feature_embeddings) # Define scale embeddings for multiscale region tokens self.num_multiscale_regions = config['parameters']['num_multiscale_regions'] self.scale_embeddings = nn.Embedding(self.num_multiscale_regions, hidden_dim) nn.init.normal_(self.scale_embeddings.weight, std=0.02) # Instantiate the prompt and region attention layers self.num_decoder_layers = config['architecture']['num_decoder_layers'] self.num_attention_heads = config['architecture']['num_attention_heads'] self.prompt_attention_layers = nn.ModuleList([ AttentionLayer(hidden_dim, hidden_dim, hidden_dim, num_heads=self.num_attention_heads) for _ in range(self.num_decoder_layers) ]) self.prompt_attention_norms = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(self.num_decoder_layers)]) self.region_attention_layers = nn.ModuleList([ CrossAttentionBlock(q_dim=hidden_dim, kv_dim=hidden_dim, hidden_dim=hidden_dim, mlp_dim=2 * hidden_dim, num_heads=self.num_attention_heads, dropout=0.1, use_bias=False) for _ in range(self.num_decoder_layers) ]) self.region_attention_norms = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(self.num_decoder_layers)]) # Instantiate the region token prediction head self.token_prediction_head = AttentionLayer(hidden_dim, hidden_dim, hidden_dim, num_heads=1, dropout=0.0, use_v_proj=False, use_out_proj=False) # Instantiate the text alignment head self.text_alignment_block = TextAlignmentBlock(hidden_dim, 2 * hidden_dim, text_embed_dim) # Instantiate the token aggregator self.token_aggregator = TokenAggregator(config) def load_state_dict_resolution_agnostic(self, state_dict, strict=False): model_state = self.state_dict() new_state = dict(state_dict) # Interpolate location_embeddings if spatial size differs if 'location_embeddings' in new_state and new_state['location_embeddings'].shape != model_state['location_embeddings'].shape: old = new_state['location_embeddings'] target_shape = model_state['location_embeddings'].shape if old.shape[0] == target_shape[0]: resized = F.interpolate(old.unsqueeze(0), size=(target_shape[1], target_shape[2]), mode='bilinear', align_corners=False) new_state['location_embeddings'] = resized.squeeze(0) else: new_state['location_embeddings'] = model_state['location_embeddings'].clone() # Interpolate feature_embeddings if spatial size differs if 'feature_embeddings' in new_state and new_state['feature_embeddings'].shape != model_state['feature_embeddings'].shape: old = new_state['feature_embeddings'] target = model_state['feature_embeddings'] if old.shape[1] == target.shape[1]: num_pos_old, C = old.shape num_pos_new = target.shape[0] h_old = int(round(num_pos_old ** 0.5)) w_old = num_pos_old // h_old h_new = int(round(num_pos_new ** 0.5)) w_new = num_pos_new // h_new old_2d = old.view(h_old, w_old, C).permute(2, 0, 1).unsqueeze(0) resized = F.interpolate(old_2d, size=(h_new, w_new), mode='bilinear', align_corners=False) new_state['feature_embeddings'] = resized.squeeze(0).permute(1, 2, 0).reshape(-1, C) else: new_state['feature_embeddings'] = model_state['feature_embeddings'].clone() return self.load_state_dict(new_state, strict=strict) def forward(self, feature_maps, grid_points, aggregate_tokens=False, remove_singleton_groups=True): if isinstance(grid_points, list): grid_points = torch.stack([gp.to(feature_maps.device) for gp in grid_points]) batch_size, num_prompts, _ = grid_points.shape # Create scale prompt embeddings for multiscale region tokens scale_prompt_embeddings = self.scale_embeddings.weight.unsqueeze(0).repeat(batch_size, 1, 1) scale_prompt_embeddings = scale_prompt_embeddings.unsqueeze(1).repeat(1, num_prompts, 1, 1) # Create spatial prompt embeddings to encode the location of the point prompts spatial_prompt_embeddings = self.location_embeddings[:, grid_points[..., 0], grid_points[..., 1]] spatial_prompt_embeddings = spatial_prompt_embeddings.permute(1, 2, 0).unsqueeze(2) spatial_prompt_embeddings = spatial_prompt_embeddings.repeat(1, 1, self.num_multiscale_regions, 1) # Create the query tokens q = scale_prompt_embeddings # Get the key and value tokens for the region attention layers kv = feature_maps.flatten(-2).permute(0, 2, 1) kv = kv + self.feature_embeddings[None] # Apply the region attention layers and the prompt attention layers for layer_idx in range(self.num_decoder_layers): q += spatial_prompt_embeddings # Apply the region attention layer q = q.reshape(batch_size, num_prompts * self.num_multiscale_regions, -1) q, _ = self.region_attention_layers[layer_idx](q, kv) q = q.reshape(batch_size, num_prompts, self.num_multiscale_regions, -1) q = self.region_attention_norms[layer_idx](q) # Apply the prompt attention layer q = q.reshape(batch_size * num_prompts, self.num_multiscale_regions, -1) q, _ = self.prompt_attention_layers[layer_idx](q, q, q) q = self.prompt_attention_norms[layer_idx](q) q = q.reshape(batch_size, num_prompts, self.num_multiscale_regions, -1) prompt_tokens = q # Get the region tokens q = prompt_tokens.reshape(batch_size, num_prompts * self.num_multiscale_regions, -1) k = kv v = kv - self.feature_embeddings[None] pred_tokens, attn_weights = self.token_prediction_head(q, k, v) pred_tokens = pred_tokens.reshape(batch_size, num_prompts, self.num_multiscale_regions, -1) attn_weights = attn_weights.reshape(batch_size, num_prompts, self.num_multiscale_regions, -1) # Get the region masks region_masks = attn_weights / attn_weights.max(dim=-1, keepdim=True)[0] region_masks = region_masks.reshape(batch_size, num_prompts, self.num_multiscale_regions, self.feature_map_resolution, self.feature_map_resolution) # Get text aligned tokens text_aligned_tokens = self.text_alignment_block(pred_tokens) outputs = { 'pred_tokens': pred_tokens, 'region_masks': region_masks, 'text_aligned_tokens': text_aligned_tokens, } if aggregate_tokens: outputs = self.token_aggregator(outputs, remove_singleton_groups=remove_singleton_groups) return outputs if __name__ == '__main__': import yaml from tqdm import tqdm from dataloader import COCOStuffDataset # Load the config with open('configs/train_dinov3_vitl16.yaml', 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) device = 'cuda' if torch.cuda.is_available() else 'cpu' # Load the dataset dataset = COCOStuffDataset(config, 'val') # Generate the grid points image_resolution = config['parameters']['image_resolution'] patch_size = config['architecture']['patch_size'] grid_size = image_resolution // patch_size x_coords = np.linspace(patch_size // 2, image_resolution - patch_size // 2, grid_size, dtype=int) y_coords = np.linspace(patch_size // 2, image_resolution - patch_size // 2, grid_size, dtype=int) grid_points = np.array([(y, x) for y in y_coords for x in x_coords]) grid_points = torch.tensor(grid_points)[None] # Get the models region_encoder = RegionEncoder(config).to(device) feature_extractor = FeatureExtractor(config, device) # Generate the region tokens for item in tqdm(dataset): image = item[0].to(device) feature_maps = feature_extractor(image[None])['feature_maps'] ren_outputs = region_encoder(feature_maps, grid_points, aggregate_tokens=True) print(ren_outputs['pred_tokens'][0].shape) print(ren_outputs['region_masks'][0].shape) print(ren_outputs['text_aligned_tokens'][0].shape) break