VisionLanguageGroup/MicroscopyMatching
0
1import os2from typing import Any, List, Optional3from huggingface_hub import hf_hub_download4from pytorch_lightning.utilities.types import STEP_OUTPUT5import torch6from PIL import Image7import numpy as np8import tifffile9from config import RunConfig10from _utils import attn_utils_new as attn_utils11from _utils.attn_utils_new import AttentionStore12from _utils.misc_helper import *13import torch.nn.functional as F14from tqdm import tqdm15import torch.nn as nn16import cv217import warnings18warnings.filterwarnings("ignore", category=UserWarning)19import pytorch_lightning as pl20from _utils.load_models import load_stable_diffusion_model21from models.model import Counting_with_SD_features_track as Counting22from models.enc_model.loca import build_model as build_loca_model23import time24from models.tra_post_model.model import TrackingTransformer25from models.tra_post_model.utils import (26 normalize,27)28from models.tra_post_model.data import build_windows_sd, get_features29from models.tra_post_model.tracking import TrackGraph, build_graph, track_greedy30import torchvision.transforms as T31from pathlib import Path32import dask.array as da33from typing import Dict, List, Optional, Union, Literal34from scipy.sparse import SparseEfficiencyWarning, csr_array35import tracemalloc36import gc37from _utils.load_track_data import load_track_images38 39SCALE = 140 41def get_instance_boxes(mask):42 # Convert to int64 if needed43 if mask.dtype != torch.long:44 mask = mask.to(torch.long)45 46 boxes = []47 instance_ids = torch.unique(mask)48 instance_ids = instance_ids[instance_ids != 0] # skip background49 50 for inst_id in instance_ids:51 inst_mask = mask == inst_id52 y_indices, x_indices = torch.where(inst_mask)53 54 if len(x_indices) == 0 or len(y_indices) == 0:55 continue56 57 x_min = torch.min(x_indices).item()58 x_max = torch.max(x_indices).item()59 y_min = torch.min(y_indices).item()60 y_max = torch.max(y_indices).item()61 62 boxes.append([x_min, y_min, x_max, y_max])63 boxes = torch.tensor(boxes, dtype=torch.float32)64 return boxes65 66class TrackingModule(pl.LightningModule):67 def __init__(self, use_box=False):68 super().__init__()69 self.use_box = use_box70 self.config = RunConfig() # config for stable diffusion71 self.initialize_model()72 73 def initialize_model(self):74 75 # load loca model76 self.loca_model = build_loca_model()77 78 self.counting_adapter = Counting(scale_factor=SCALE)79 80 ### load stable diffusion and its controller81 self.stable = load_stable_diffusion_model(config=self.config)82 self.noise_scheduler = self.stable.scheduler83 self.controller = AttentionStore(max_size=64)84 attn_utils.register_attention_control(self.stable, self.controller)85 attn_utils.register_hier_output(self.stable)86 87 ##### initialize token_emb #####88 placeholder_token = "<task-prompt>"89 self.task_token = "repetitive objects"90 # Add the placeholder token in tokenizer91 num_added_tokens = self.stable.tokenizer.add_tokens(placeholder_token)92 if num_added_tokens == 0:93 raise ValueError(94 f"The tokenizer already contains the token {placeholder_token}. Please pass a different"95 " `placeholder_token` that is not already in the tokenizer."96 )97 try:98 task_embed_from_pretrain = hf_hub_download(99 repo_id="phoebe777777/111",100 filename="task_embed.pth",101 token=None,102 force_download=False103 )104 placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(placeholder_token)105 self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))106 107 token_embeds = self.stable.text_encoder.get_input_embeddings().weight.data108 token_embeds[placeholder_token_id] = task_embed_from_pretrain109 except:110 #@title Get token ids for our placeholder and initializer token. This code block will complain if initializer string is not a single token111 # Convert the initializer_token, placeholder_token to ids112 initializer_token = "track"113 token_ids = self.stable.tokenizer.encode(initializer_token, add_special_tokens=False)114 # Check if initializer_token is a single token or a sequence of tokens115 if len(token_ids) > 1:116 raise ValueError("The initializer token must be a single token.")117 118 initializer_token_id = token_ids[0]119 placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(placeholder_token)120 121 self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))122 123 token_embeds = self.stable.text_encoder.get_input_embeddings().weight.data124 token_embeds[placeholder_token_id] = token_embeds[initializer_token_id]125 126 # others127 self.placeholder_token = placeholder_token128 self.placeholder_token_id = placeholder_token_id129 130 fpath = Path("_utils/config.yaml")131 132 model = TrackingTransformer.from_cfg(133 cfg_path=fpath,134 )135 136 self.track_model = model137 138 139 def move_to_device(self, device):140 self.stable.to(device)141 self.counting_adapter.to(device)142 self.loca_model.to(device)143 self.track_model.to(device)144 145 self.to(device)146 147 def on_train_start(self) -> None:148 device = self.device149 dtype = self.dtype150 self.stable.to(device,dtype)151 152 def on_validation_start(self) -> None:153 device = self.device154 dtype = self.dtype155 self.stable.to(device,dtype)156 157 def forward(self, data):158 159 input_image_stable = data["image_stable"]160 boxes = data["boxes"]161 input_image = data["img_enc"]162 mask = data["mask"]163 latents = self.stable.vae.encode(input_image_stable).latent_dist.sample().detach()164 latents = latents * 0.18215165 # Sample noise that we'll add to the latents166 noise = torch.randn_like(latents)167 bsz = latents.shape[0]168 timesteps = torch.tensor([20], device=latents.device).long()169 noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)170 input_ids_ = self.stable.tokenizer(171 self.placeholder_token,172 # "object",173 padding="max_length",174 truncation=True,175 max_length=self.stable.tokenizer.model_max_length,176 return_tensors="pt",177 )178 input_ids = input_ids_["input_ids"].to(self.device)179 attention_mask = input_ids_["attention_mask"].to(self.device)180 encoder_hidden_states = self.stable.text_encoder(input_ids, attention_mask)[0]181 encoder_hidden_states = encoder_hidden_states.repeat(bsz, 1, 1)182 183 time1 = time.time()184 input_image = input_image.to(self.device)185 boxes = boxes.to(self.device)186 187 loca_out = self.loca_model.forward_before_reg(input_image, boxes)188 loca_feature_bf_regression = loca_out["feature_bf_regression"]189 # time2 = time.time()190 191 task_loc_idx = torch.nonzero(input_ids == self.placeholder_token_id)192 adapted_emb = self.counting_adapter.adapter(loca_feature_bf_regression, boxes) # shape [1, 768]193 194 if task_loc_idx.shape[0] == 0:195 encoder_hidden_states[0,2,:] = adapted_emb.squeeze()196 else:197 encoder_hidden_states[:,task_loc_idx[0, 1]+1,:] = adapted_emb.squeeze() 198 199 # Predict the noise residual200 noise_pred, feature_list = self.stable.unet(noisy_latents, timesteps, encoder_hidden_states)201 time3 = time.time()202 noise_pred = noise_pred.sample203 204 attention_store = self.controller.attention_store205 206 # print(time2-time1, time3-time2)207 208 attention_maps = []209 exemplar_attention_maps = []210 211 cross_self_task_attn_maps = []212 cross_self_exe_attn_maps = []213 214 # only use 64x64 self-attention215 self_attn_aggregate = attn_utils.aggregate_attention( # [res, res, 4096]216 prompts=[self.config.prompt for i in range(bsz)], 217 attention_store=self.controller, 218 res=64,219 from_where=("up", "down"),220 is_cross=False,221 select=0222 )223 self_attn_aggregate32 = attn_utils.aggregate_attention( # [res, res, 4096]224 prompts=[self.config.prompt for i in range(bsz)], 225 attention_store=self.controller, 226 res=32,227 from_where=("up", "down"),228 is_cross=False,229 select=0230 )231 self_attn_aggregate16 = attn_utils.aggregate_attention( # [res, res, 4096]232 prompts=[self.config.prompt for i in range(bsz)], 233 attention_store=self.controller, 234 res=16,235 from_where=("up", "down"),236 is_cross=False,237 select=0238 )239 240 # cross attention241 for res in [32, 16]:242 attn_aggregate = attn_utils.aggregate_attention( # [res, res, 77]243 prompts=[self.config.prompt for i in range(bsz)], 244 attention_store=self.controller, 245 res=res,246 from_where=("up", "down"),247 is_cross=True,248 select=0249 )250 251 task_attn_ = attn_aggregate[:, :, 1].unsqueeze(0).unsqueeze(0) # [1, 1, res, res]252 attention_maps.append(task_attn_)253 exemplar_attns = attn_aggregate[:, :, 2].unsqueeze(0).unsqueeze(0) 254 exemplar_attention_maps.append(exemplar_attns)255 256 257 scale_factors = [(64 // attention_maps[i].shape[-1]) for i in range(len(attention_maps))]258 attns = torch.cat([F.interpolate(attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(attention_maps))])259 task_attn_64 = torch.mean(attns, dim=0, keepdim=True)260 261 262 scale_factors = [(64 // exemplar_attention_maps[i].shape[-1]) for i in range(len(exemplar_attention_maps))]263 attns = torch.cat([F.interpolate(exemplar_attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps))])264 exemplar_attn_64 = torch.mean(attns, dim=0, keepdim=True)265 266 cross_self_task_attn = attn_utils.self_cross_attn(self_attn_aggregate, task_attn_64)267 cross_self_exe_attn = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64)268 cross_self_task_attn_maps.append(cross_self_task_attn)269 cross_self_exe_attn_maps.append(cross_self_exe_attn)270 271 task_attn_64 = (task_attn_64 - task_attn_64.min()) / (task_attn_64.max() - task_attn_64.min() + 1e-6)272 cross_self_task_attn = (cross_self_task_attn - cross_self_task_attn.min()) / (cross_self_task_attn.max() - cross_self_task_attn.min() + 1e-6)273 exemplar_attn_64 = (exemplar_attn_64 - exemplar_attn_64.min()) / (exemplar_attn_64.max() - exemplar_attn_64.min() + 1e-6)274 cross_self_exe_attn = (cross_self_exe_attn - cross_self_exe_attn.min()) / (cross_self_exe_attn.max() - cross_self_exe_attn.min() + 1e-6)275 276 attn_stack = [task_attn_64 / 2, cross_self_task_attn / 2, exemplar_attn_64, cross_self_exe_attn]277 attn_stack = torch.cat(attn_stack, dim=1)278 279 280 attn_after_new_regressor, loss = self.counting_adapter.regressor(input_image, attn_stack, feature_list, mask.cpu().numpy(), training=False) 281 282 return {283 "attn_after_new_regressor":attn_after_new_regressor, 284 "task_attn_64":task_attn_64, 285 "cross_self_task_attn":cross_self_task_attn, 286 "exemplar_attn_64": exemplar_attn_64,287 "cross_self_exe_attn": cross_self_exe_attn,288 "noise_pred":noise_pred,289 "noise":noise,290 "self_attn_aggregate":self_attn_aggregate,291 "self_attn_aggregate32":self_attn_aggregate32,292 "self_attn_aggregate16":self_attn_aggregate16,293 "loss": loss294 }295 296 def forward_sd(self, input_image_stable, input_image, boxes, height, width, mask=None):297 298 input_image_stable = input_image_stable.to(self.device)299 # density = data["density"]300 if boxes is not None:301 boxes = boxes.to(self.device)302 input_image = input_image.to(self.device)303 if mask is not None:304 mask = mask.to(self.device)305 else:306 mask = torch.zeros((input_image.shape[0], 1, input_image.shape[2], input_image.shape[3])).to(self.device)307 308 latents = self.stable.vae.encode(input_image_stable).latent_dist.sample().detach()309 latents = latents * 0.18215310 # Sample noise that we'll add to the latents311 noise = torch.randn_like(latents)312 bsz = latents.shape[0]313 timesteps = torch.tensor([20], device=latents.device).long()314 noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)315 input_ids_ = self.stable.tokenizer(316 self.placeholder_token + " " + self.task_token,317 # "object",318 padding="max_length",319 truncation=True,320 max_length=self.stable.tokenizer.model_max_length,321 return_tensors="pt",322 )323 input_ids = input_ids_["input_ids"].to(self.device)324 attention_mask = input_ids_["attention_mask"].to(self.device)325 encoder_hidden_states = self.stable.text_encoder(input_ids, attention_mask)[0]326 encoder_hidden_states = encoder_hidden_states.repeat(bsz, 1, 1)327 328 329 if boxes is not None and not self.training:330 if self.adapt_emb is None:331 loca_out_ = self.loca_model.forward_before_reg(input_image, boxes)332 loca_feature_bf_regression_ = loca_out_["feature_bf_regression"]333 adapted_emb = self.counting_adapter.adapter(loca_feature_bf_regression_, boxes) # shape [1, 768]334 else:335 adapted_emb = self.adapt_emb.to(self.device)336 task_loc_idx = torch.nonzero(input_ids == self.placeholder_token_id)337 if task_loc_idx.shape[0] == 0:338 encoder_hidden_states[0,5,:] = adapted_emb.squeeze() 339 else:340 encoder_hidden_states[:,task_loc_idx[0, 1]+4,:] = adapted_emb.squeeze() 341 342 # Predict the noise residual343 noise_pred, feature_list = self.stable.unet(noisy_latents, timesteps, encoder_hidden_states)344 noise_pred = noise_pred.sample345 attention_store = self.controller.attention_store346 347 348 attention_maps = []349 exemplar_attention_maps = []350 exemplar_attention_maps1 = []351 exemplar_attention_maps2 = []352 exemplar_attention_maps3 = []353 exemplar_attention_maps4 = []354 355 cross_self_task_attn_maps = []356 cross_self_exe_attn_maps = []357 358 # only use 64x64 self-attention359 self_attn_aggregate = attn_utils.aggregate_attention( # [res, res, 4096]360 prompts=[self.config.prompt for i in range(bsz)], 361 attention_store=self.controller, 362 res=64,363 from_where=("up", "down"),364 is_cross=False,365 select=0366 )367 368 # cross attention369 for res in [32, 16]:370 attn_aggregate = attn_utils.aggregate_attention( # [res, res, 77]371 prompts=[self.config.prompt for i in range(bsz)], 372 attention_store=self.controller, 373 res=res,374 from_where=("up", "down"),375 is_cross=True,376 select=0377 )378 379 task_attn_ = attn_aggregate[:, :, 1].unsqueeze(0).unsqueeze(0) # [1, 1, res, res]380 attention_maps.append(task_attn_)381 # if self.boxes is not None and not self.training:382 exemplar_attns1 = attn_aggregate[:, :, 2].unsqueeze(0).unsqueeze(0) 383 exemplar_attention_maps1.append(exemplar_attns1)384 exemplar_attns2 = attn_aggregate[:, :, 3].unsqueeze(0).unsqueeze(0) 385 exemplar_attention_maps2.append(exemplar_attns2)386 exemplar_attns3 = attn_aggregate[:, :, 4].unsqueeze(0).unsqueeze(0) 387 exemplar_attention_maps3.append(exemplar_attns3)388 exemplar_attns4 = attn_aggregate[:, :, 5].unsqueeze(0).unsqueeze(0) 389 exemplar_attention_maps4.append(exemplar_attns4)390 391 392 393 scale_factors = [(64 // attention_maps[i].shape[-1]) for i in range(len(attention_maps))]394 attns = torch.cat([F.interpolate(attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(attention_maps))])395 task_attn_64 = torch.mean(attns, dim=0, keepdim=True)396 cross_self_task_attn = attn_utils.self_cross_attn(self_attn_aggregate, task_attn_64)397 cross_self_task_attn_maps.append(cross_self_task_attn)398 399 # if not self.training:400 scale_factors = [(64 // exemplar_attention_maps1[i].shape[-1]) for i in range(len(exemplar_attention_maps1))]401 attns = torch.cat([F.interpolate(exemplar_attention_maps1[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps1))])402 exemplar_attn_64_1 = torch.mean(attns, dim=0, keepdim=True)403 404 scale_factors = [(64 // exemplar_attention_maps2[i].shape[-1]) for i in range(len(exemplar_attention_maps2))]405 attns = torch.cat([F.interpolate(exemplar_attention_maps2[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps2))])406 exemplar_attn_64_2 = torch.mean(attns, dim=0, keepdim=True)407 408 scale_factors = [(64 // exemplar_attention_maps3[i].shape[-1]) for i in range(len(exemplar_attention_maps3))]409 attns = torch.cat([F.interpolate(exemplar_attention_maps3[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps3))])410 exemplar_attn_64_3 = torch.mean(attns, dim=0, keepdim=True)411 412 if boxes is not None:413 scale_factors = [(64 // exemplar_attention_maps4[i].shape[-1]) for i in range(len(exemplar_attention_maps4))]414 attns = torch.cat([F.interpolate(exemplar_attention_maps4[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps4))])415 exemplar_attn_64_4 = torch.mean(attns, dim=0, keepdim=True)416 417 exes = []418 cross_exes = []419 cross_self_exe_attn1 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_1)420 cross_self_exe_attn2 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_2)421 cross_self_exe_attn3 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_3)422 423 # # average424 exemplar_attn_64_1 = (exemplar_attn_64_1 - exemplar_attn_64_1.min()) / (exemplar_attn_64_1.max() - exemplar_attn_64_1.min() + 1e-6)425 exemplar_attn_64_2 = (exemplar_attn_64_2 - exemplar_attn_64_2.min()) / (exemplar_attn_64_2.max() - exemplar_attn_64_2.min() + 1e-6)426 exemplar_attn_64_3 = (exemplar_attn_64_3 - exemplar_attn_64_3.min()) / (exemplar_attn_64_3.max() - exemplar_attn_64_3.min() + 1e-6)427 cross_self_exe_attn1 = (cross_self_exe_attn1 - cross_self_exe_attn1.min()) / (cross_self_exe_attn1.max() - cross_self_exe_attn1.min() + 1e-6)428 cross_self_exe_attn2 = (cross_self_exe_attn2 - cross_self_exe_attn2.min()) / (cross_self_exe_attn2.max() - cross_self_exe_attn2.min() + 1e-6)429 cross_self_exe_attn3 = (cross_self_exe_attn3 - cross_self_exe_attn3.min()) / (cross_self_exe_attn3.max() - cross_self_exe_attn3.min() + 1e-6)430 exes = [exemplar_attn_64_1, exemplar_attn_64_2, exemplar_attn_64_3]431 cross_exes = [cross_self_exe_attn1, cross_self_exe_attn2, cross_self_exe_attn3]432 if boxes is not None:433 cross_self_exe_attn4 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_4)434 exemplar_attn_64_4 = (exemplar_attn_64_4 - exemplar_attn_64_4.min()) / (exemplar_attn_64_4.max() - exemplar_attn_64_4.min() + 1e-6)435 cross_self_exe_attn4 = (cross_self_exe_attn4 - cross_self_exe_attn4.min()) / (cross_self_exe_attn4.max() - cross_self_exe_attn4.min() + 1e-6)436 exes.append(exemplar_attn_64_4)437 cross_exes.append(cross_self_exe_attn4)438 exemplar_attn_64 = sum(exes) / len(exes)439 cross_self_exe_attn = sum(cross_exes) / len(cross_exes)440 441 442 443 if self.use_box:444 attn_stack = [task_attn_64 / 2, cross_self_task_attn / 2, exemplar_attn_64, cross_self_exe_attn]445 else:446 attn_stack = [exemplar_attn_64 / 2, cross_self_exe_attn / 2, exemplar_attn_64, cross_self_exe_attn]447 attn_stack = torch.cat(attn_stack, dim=1)448 449 450 attn_after_new_regressor, loss, _ = self.counting_adapter.regressor.forward_seg(input_image, attn_stack, feature_list, mask.cpu().numpy(), self.training) 451 452 if not self.training:453 pred_mask = attn_after_new_regressor.detach().cpu()454 pred_boxes = get_instance_boxes(pred_mask.squeeze())455 456 self.boxes = pred_boxes.unsqueeze(0)457 458 if pred_boxes.shape[0] == 0:459 print("No instances detected in the predicted mask.")460 self.adapt_emb = adapted_emb.detach().cpu() # reuse emb461 else:462 pred_boxes = pred_boxes.unsqueeze(0).to(self.device)463 loca_out_ = self.loca_model.forward_before_reg(input_image, pred_boxes)464 loca_feature_bf_regression_ = loca_out_["feature_bf_regression"]465 adapted_emb_ = self.counting_adapter.adapter(loca_feature_bf_regression_, pred_boxes) # shape [1, 768]466 self.adapt_emb = adapted_emb_.detach().cpu()467 468 # resize to original image size469 mask_np = attn_after_new_regressor.squeeze().detach().cpu().numpy()470 mask_resized = cv2.resize(mask_np, (width, height), interpolation=cv2.INTER_NEAREST)471 472 return mask_resized473 474 def forward_boxes(self, input_image_stable, boxes, input_image):475 476 latents = self.stable.vae.encode(input_image_stable).latent_dist.sample().detach()477 latents = latents * 0.18215478 # Sample noise that we'll add to the latents479 noise = torch.randn_like(latents)480 bsz = latents.shape[0]481 timesteps = torch.tensor([20], device=latents.device).long()482 noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)483 input_ids_ = self.stable.tokenizer(484 self.placeholder_token,485 # "object",486 padding="max_length",487 truncation=True,488 max_length=self.stable.tokenizer.model_max_length,489 return_tensors="pt",490 )491 input_ids = input_ids_["input_ids"].to(self.device)492 attention_mask = input_ids_["attention_mask"].to(self.device)493 encoder_hidden_states = self.stable.text_encoder(input_ids, attention_mask)[0]494 encoder_hidden_states = encoder_hidden_states.repeat(bsz, 1, 1)495 496 time1 = time.time()497 input_image = input_image.to(self.device)498 boxes = boxes.to(self.device)499 500 loca_out = self.loca_model.forward_before_reg(input_image, boxes)501 loca_feature_bf_regression = loca_out["feature_bf_regression"]502 # time2 = time.time()503 504 task_loc_idx = torch.nonzero(input_ids == self.placeholder_token_id)505 adapted_emb = self.counting_adapter.adapter.forward_boxes(loca_feature_bf_regression, boxes) # shape [n_instance, 768]506 n_instance = adapted_emb.shape[0]507 n_forward = int(np.ceil(n_instance / 74)) # in total 75 prompts including 1 task prompt and 74 object prompts?508 509 task_cross_attention = []510 instances_cross_attention = []511 512 for n in range(n_forward):513 len_ = min(74, n_instance - n * 74)514 encoder_hidden_states[:,(task_loc_idx[0, 1]+1):(task_loc_idx[0, 1]+1+len_),:] = adapted_emb[n*74:n*74+len_].squeeze() 515 516 517 # Predict the noise residual518 noise_pred, feature_list = self.stable.unet(noisy_latents, timesteps, encoder_hidden_states)519 noise_pred = noise_pred.sample520 521 522 523 attention_maps = []524 exemplar_attention_maps = []525 526 # cross attention527 for res in [32, 16]:528 attn_aggregate = attn_utils.aggregate_attention( # [res, res, 77]529 prompts=[self.config.prompt for i in range(bsz)], 530 attention_store=self.controller, 531 res=res,532 from_where=("up", "down"),533 is_cross=True,534 select=0535 )536 537 task_attn_ = attn_aggregate[:, :, 1].unsqueeze(0).unsqueeze(0) # [1, 1, res, res]538 attention_maps.append(task_attn_)539 try:540 exemplar_attns = attn_aggregate[:, :, (task_loc_idx[0, 1]+1):(task_loc_idx[0, 1]+1+len_)].unsqueeze(0) 541 except:542 print(n_instance, len_)543 exemplar_attns = torch.permute(exemplar_attns, (0, 3, 1, 2)) # [1, len_, res, res]544 exemplar_attention_maps.append(exemplar_attns)545 546 547 548 scale_factors = [(64 // attention_maps[i].shape[-1]) for i in range(len(attention_maps))]549 attns = torch.cat([F.interpolate(attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(attention_maps))])550 task_attn_64 = torch.mean(attns, dim=0, keepdim=True)551 552 try:553 scale_factors = [(64 // exemplar_attention_maps[i].shape[-1]) for i in range(len(exemplar_attention_maps))]554 attns = torch.cat([F.interpolate(exemplar_attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps))])555 except:556 print("exemplar_attention_maps shape mismatch, n_instance: {}, len_: {}".format(n_instance, len_))557 print(exemplar_attention_maps[0].shape)558 print(exemplar_attention_maps[1].shape)559 print(exemplar_attention_maps[2].shape)560 exemplar_attn_64 = torch.mean(attns, dim=0, keepdim=True)561 562 563 task_attn_64 = (task_attn_64 - task_attn_64.min()) / (task_attn_64.max() - task_attn_64.min() + 1e-6)564 exemplar_attn_64 = (exemplar_attn_64 - exemplar_attn_64.min()) / (exemplar_attn_64.max() - exemplar_attn_64.min() + 1e-6)565 566 task_cross_attention.append(task_attn_64)567 instances_cross_attention.append(exemplar_attn_64)568 569 task_cross_attention = torch.cat(task_cross_attention, dim=0) # [n_forward, 1, 64, 64]570 task_cross_attention = torch.mean(task_cross_attention, dim=0, keepdim=True) # [1, 1, 64, 64]571 instances_cross_attention = torch.cat(instances_cross_attention, dim=1) # [1, n_instance, 64, 64]572 assert instances_cross_attention.shape[1] == n_instance, "instances_cross_attention shape mismatch"573 attn_stack = [task_cross_attention / 2, instances_cross_attention]574 attn_stack = torch.cat(attn_stack, dim=1)575 576 del exemplar_attention_maps, attention_maps, attns, task_attn_64, exemplar_attn_64, latents577 del input_ids_, input_ids, attention_mask, encoder_hidden_states, timesteps, noisy_latents578 del loca_out, loca_feature_bf_regression, adapted_emb579 torch.cuda.empty_cache()580 581 return {582 "task_attn_64":task_cross_attention, 583 "exemplar_attn_64": instances_cross_attention,584 "noise_pred":noise_pred,585 "noise":noise,586 "attn_stack": attn_stack,587 "feature_list": feature_list,588 }589 590 591 def common_step(self, batch):592 mask = batch["mask_t"].to(torch.float32).to(self.device)593 if mask.dim() == 3:594 mask = mask.unsqueeze(0)595 596 image_stable = batch["image_stable"]597 boxes = batch["boxes"]598 input_image = batch["img_enc"]599 input_image = input_image.to(self.device)600 image_stable = image_stable.to(self.device)601 keep_boxes = None602 if image_stable.dim() == 4:603 image_stable = image_stable.unsqueeze(0)604 if input_image.dim() == 4:605 input_image = input_image.unsqueeze(0)606 607 608 # segmentation part609 n_frames = mask.shape[1]610 masks_pred = []611 612 for i in range(n_frames):613 mask_ = mask[:, i, :, :].unsqueeze(0) # [1, 1, H, W]614 mask_ = F.interpolate(mask_.float(), size=(512, 512), mode='nearest') # [1, 1, 512, 512]615 mask_ = mask_.to(torch.int64).squeeze(0).detach().to(self.device) # [1, 512, 512]616 masks_pred.append(mask_)617 del mask_618 619 620 621 # if True:622 attns_emb = []623 for i in range(n_frames):624 image_stable_prev = image_stable[:, max(0, i-1), :, :, :]625 image_stable_after = image_stable[:, min(n_frames-1, i+1), :, :, :]626 input_image_curr = input_image[:, i, :, :, :]627 628 mask_ = masks_pred[i].detach()629 unique_labels = torch.unique(mask_) # tensor([0, 1, 2, ...])630 boxes_all = []631 for label in unique_labels:632 if label.item() == 0:633 continue634 binary_mask = (mask_[0] == label).to(torch.uint8) # [H, W]635 636 # 找非零点坐标637 y_coords, x_coords = torch.nonzero(binary_mask, as_tuple=True)638 if len(x_coords) == 0 or len(y_coords) == 0:639 continue640 x_min = torch.min(x_coords)641 y_min = torch.min(y_coords)642 x_max = torch.max(x_coords)643 y_max = torch.max(y_coords)644 boxes_all.append([x_min.item(), y_min.item(), x_max.item(), y_max.item()])645 boxes_all_t = torch.tensor(boxes_all, dtype=torch.float32).unsqueeze(0)646 647 648 output_prev = self.forward_boxes(image_stable_prev, boxes_all_t, input_image_curr)649 attn_prev = output_prev["exemplar_attn_64"]650 feature_list_prev = output_prev["feature_list"]651 652 output_after = self.forward_boxes(image_stable_after, boxes_all_t, input_image_curr)653 # attn_stack = output["attn_stack"]654 attn_after = output_after["exemplar_attn_64"] # [1, n_instance, 64, 64]655 feature_list_after = output_after["feature_list"] # [1, n_channels, res, res]656 657 attn_prev = torch.permute(attn_prev, (1, 0, 2, 3)) # [n_instance, 1, 64, 64]658 attn_after = torch.permute(attn_after, (1, 0, 2, 3))659 attn_emb = self.counting_adapter.regressor(attn_prev, feature_list_prev, attn_after, feature_list_after)660 attns_emb.append(attn_emb.detach())661 662 attns_emb = torch.cat(attns_emb, dim=1) # [1, n_instance, 4]663 # tracking part664 665 feats = batch["features_t"]666 coords = batch["coords_t"]667 668 with torch.no_grad():669 670 A_pred = self.track_model(coords, feats, attn_feat=attns_emb).detach()671 672 del masks_pred, feats, coords, batch673 gc.collect()674 torch.cuda.empty_cache()675 torch.cuda.ipc_collect()676 677 return A_pred678 679 # @profile680 def _predict_batch(self, batch):681 feats = batch["features_t"].to(self.device)682 coords = batch["coords_t"].to(self.device)683 timepoints = batch["timepoints_t"].to(self.device)684 # Hack that assumes that all parameters of a model are on the same device685 device = next(self.track_model.parameters()).device686 feats = feats.unsqueeze(0).to(device)687 timepoints = timepoints.unsqueeze(0).to(device)688 coords = coords.unsqueeze(0).to(device)689 690 # Concat timepoints to coordinates691 coords = torch.cat((timepoints.unsqueeze(2).float(), coords), dim=2)692 batch["coords_t"] = coords693 batch["features_t"] = feats694 with torch.no_grad():695 A = self.common_step(batch)696 torch.cuda.empty_cache()697 gc.collect()698 699 A = self.track_model.normalize_output(A, timepoints, coords)700 701 A = A.squeeze(0).detach().cpu().numpy()702 703 del feats, coords, timepoints, batch704 705 return A706 707 # @profile708 def predict_windows(self,709 windows: List[dict],710 features: list,711 model,712 imgs_enc: Optional[np.ndarray] = None,713 imgs_stable: Optional[np.ndarray] = None,714 intra_window_weight: float = 0,715 delta_t: int = 1,716 edge_threshold: float = 0.05,717 spatial_dim: int = 3,718 progbar_class=tqdm,719 ) -> dict:720 721 # first get all objects/coords722 time_labels_to_id = dict()723 node_properties = list()724 max_id = np.sum([len(f.labels) for f in features])725 726 all_timepoints = np.concatenate([f.timepoints for f in features])727 all_labels = np.concatenate([f.labels for f in features])728 all_coords = np.concatenate([f.coords for f in features])729 all_coords = all_coords[:, -spatial_dim:]730 731 for i, (t, la, c) in enumerate(zip(all_timepoints, all_labels, all_coords)):732 time_labels_to_id[(t, la)] = i733 node_properties.append(734 dict(735 id=i,736 coords=tuple(c),737 time=t,738 # index=ix,739 label=la,740 )741 )742 743 # create assoc matrix between ids744 sp_weights, sp_accum = (745 csr_array((max_id, max_id), dtype=np.float32),746 csr_array((max_id, max_id), dtype=np.float32),747 )748 749 tracemalloc.start()750 751 for t in progbar_class(752 range(len(windows)),753 desc="Computing associations",754 ):755 # This assumes that the samples in the dataset are ordered by time and start at 0.756 batch = windows[t]757 timepoints = batch["timepoints"]758 labels = batch["labels"]759 760 A = self._predict_batch(batch)761 762 dt = timepoints[None, :] - timepoints[:, None]763 time_mask = np.logical_and(dt <= delta_t, dt > 0)764 A[~time_mask] = 0765 ii, jj = np.where(A >= edge_threshold)766 767 if len(ii) == 0:768 continue769 770 labels_ii = labels[ii]771 labels_jj = labels[jj]772 ts_ii = timepoints[ii]773 ts_jj = timepoints[jj]774 nodes_ii = np.array(775 tuple(time_labels_to_id[(t, lab)] for t, lab in zip(ts_ii, labels_ii))776 )777 nodes_jj = np.array(778 tuple(time_labels_to_id[(t, lab)] for t, lab in zip(ts_jj, labels_jj))779 )780 781 # weight middle parts higher782 t_middle = t + (model.config["window"] - 1) / 2783 ddt = timepoints[:, None] - t_middle * np.ones_like(dt)784 window_weight = np.exp(-intra_window_weight * ddt**2) # default is 1785 # window_weight = np.exp(4*A) # smooth max786 sp_weights[nodes_ii, nodes_jj] += window_weight[ii, jj] * A[ii, jj]787 sp_accum[nodes_ii, nodes_jj] += window_weight[ii, jj]788 789 790 del batch, A, ii, jj, labels_ii, labels_jj, ts_ii, ts_jj, nodes_ii, nodes_jj, dt, time_mask791 gc.collect()792 torch.cuda.empty_cache()793 torch.cuda.ipc_collect()794 795 sp_weights_coo = sp_weights.tocoo()796 sp_accum_coo = sp_accum.tocoo()797 assert np.allclose(sp_weights_coo.col, sp_accum_coo.col) and np.allclose(798 sp_weights_coo.row, sp_accum_coo.row799 )800 801 # Normalize weights by the number of times they were written from different sliding window positions802 weights = tuple(803 ((i, j), v / a)804 for i, j, v, a in zip(805 sp_weights_coo.row,806 sp_weights_coo.col,807 sp_weights_coo.data,808 sp_accum_coo.data,809 )810 )811 812 results = dict()813 results["nodes"] = node_properties814 results["weights"] = weights815 816 return results817 818 819 def _predict(820 self,821 imgs: Union[np.ndarray, da.Array],822 masks: Union[np.ndarray, da.Array],823 imgs_enc: Optional[np.ndarray] = None,824 imgs_stable: Optional[np.ndarray] = None,825 boxes: Optional[np.ndarray] = None,826 edge_threshold: float = 0.05,827 n_workers: int = 0,828 normalize_imgs: bool = True,829 progbar_class=tqdm,830 ):831 print("Predicting weights for candidate graph")832 if normalize_imgs:833 if isinstance(imgs, da.Array):834 imgs = imgs.map_blocks(normalize)835 else:836 imgs = normalize(imgs)837 838 self.eval()839 840 features = get_features(841 detections=masks,842 imgs=imgs,843 ndim=self.track_model.config["coord_dim"],844 n_workers=n_workers,845 progbar_class=progbar_class,846 )847 print("Building windows")848 windows = build_windows_sd(849 features,850 imgs_enc=imgs_enc,851 imgs_stable=imgs_stable,852 boxes=boxes,853 imgs=imgs,854 masks=masks,855 window_size=self.track_model.config["window"],856 progbar_class=progbar_class,857 )858 859 print("Predicting windows")860 with torch.no_grad():861 predictions = self.predict_windows(862 windows=windows,863 features=features,864 imgs_enc=imgs_enc,865 imgs_stable=imgs_stable,866 model=self.track_model,867 edge_threshold=edge_threshold,868 spatial_dim=masks.ndim - 1,869 progbar_class=progbar_class,870 )871 872 return predictions873 874 def _track_from_predictions(875 self,876 predictions,877 mode: Literal["greedy_nodiv", "greedy", "ilp"] = "greedy",878 use_distance: bool = False,879 max_distance: int = 256,880 max_neighbors: int = 10,881 delta_t: int = 1,882 **kwargs,883 ):884 print("Running greedy tracker")885 nodes = predictions["nodes"]886 weights = predictions["weights"]887 888 candidate_graph = build_graph(889 nodes=nodes,890 weights=weights,891 use_distance=use_distance,892 max_distance=max_distance,893 max_neighbors=max_neighbors,894 delta_t=delta_t,895 )896 if mode == "greedy":897 return track_greedy(candidate_graph)898 elif mode == "greedy_nodiv":899 return track_greedy(candidate_graph, allow_divisions=False)900 elif mode == "ilp":901 from models.tra_post_model.tracking.ilp import track_ilp902 903 return track_ilp(candidate_graph, ilp_config="gt", **kwargs)904 else:905 raise ValueError(f"Tracking mode {mode} does not exist.")906 907 def track(908 self,909 file_dir: str,910 boxes: Optional[torch.Tensor] = None,911 mode: Literal["greedy_nodiv", "greedy", "ilp"] = "greedy",912 normalize_imgs: bool = True,913 progbar_class=tqdm,914 n_workers: int = 0,915 dataname: Optional[str] = None,916 **kwargs,917 ) -> TrackGraph:918 """Track objects across time frames.919 920 This method links segmented objects across time frames using the specified921 tracking mode. No hyperparameters need to be chosen beyond the tracking mode.922 923 Args:924 imgs: Input images of shape (T,(Z),Y,X) (numpy or dask array)925 masks: Instance segmentation masks of shape (T,(Z),Y,X).926 mode: Tracking mode:927 - "greedy_nodiv": Fast greedy linking without division928 - "greedy": Fast greedy linking with division929 - "ilp": Integer Linear Programming based linking (more accurate but slower)930 progbar_class: Progress bar class to use.931 n_workers: Number of worker processes for feature extraction.932 normalize_imgs: Whether to normalize the images.933 **kwargs: Additional arguments passed to tracking algorithm.934 935 Returns:936 TrackGraph containing the tracking results.937 """938 939 self.eval()940 imgs, imgs_raw, images_stable, tra_imgs, imgs_01, height, width = load_track_images(file_dir)941 if boxes is not None:942 boxes = torch.as_tensor(boxes, dtype=torch.float32)943 if boxes.ndim == 2:944 boxes = boxes.unsqueeze(0)945 boxes = boxes * torch.tensor(946 [512.0 / width, 512.0 / height, 512.0 / width, 512.0 / height],947 dtype=torch.float32,948 )949 boxes[..., 0::2].clamp_(0, 512)950 boxes[..., 1::2].clamp_(0, 512)951 imgs_stable = torch.from_numpy(images_stable).float().to(self.device)952 imgs_enc = torch.from_numpy(imgs).float().to(self.device)953 954 955 """get segmentation masks first"""956 self.boxes = None957 self.adapt_emb = None958 masks = []959 for i, (input_image, input_image_stable) in tqdm(enumerate(zip(imgs_enc, imgs_stable))):960 input_image = input_image.unsqueeze(0)961 input_image_stable = input_image_stable.unsqueeze(0)962 if i == 0:963 if self.use_box and boxes is not None:964 self.boxes = boxes.to(self.device)965 else:966 self.boxes = None967 968 with torch.no_grad():969 mask = self.forward_sd(input_image_stable, input_image, self.boxes, height=height, width=width)970 masks.append(mask)971 972 masks = np.stack(masks, axis=0) # (T, H, W)973 974# -------------------------975 if not masks.shape == tra_imgs.shape:976 raise RuntimeError(977 f"Img shape {tra_imgs.shape} and mask shape {masks.shape} do not match."978 )979 980 if not tra_imgs.ndim == self.track_model.config["coord_dim"] + 1:981 raise RuntimeError(982 f"images should be a sequence of {self.track_model.config['coord_dim']}D images"983 )984 985 predictions = self._predict(986 tra_imgs,987 masks,988 imgs_enc=imgs_enc,989 imgs_stable=imgs_stable,990 boxes=boxes,991 normalize_imgs=normalize_imgs,992 progbar_class=progbar_class,993 n_workers=n_workers,994 )995 track_graph = self._track_from_predictions(predictions, mode=mode, **kwargs)996 997 return track_graph, masks998