VisionLanguageGroup/MicroscopyMatching
0
1import os2from typing import Any, List, Optional3from huggingface_hub import hf_hub_download4from pytorch_lightning.utilities.types import STEP_OUTPUT5import torch6import os7from PIL import Image8import numpy as np9from 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 F14import logging15import matplotlib.pyplot as plt16import matplotlib.patches as patches17import cv218import warnings19warnings.filterwarnings("ignore", category=UserWarning)20import pytorch_lightning as pl21from _utils.load_models import load_stable_diffusion_model22from models.model import Counting_with_SD_features_dino_vit_c3 as Counting23from models.enc_model.loca import build_model as build_loca_model24import time25from models.seg_post_model import metrics26from datetime import datetime27import json28import logging29from PIL import Image30import torchvision.transforms as T31import cv232from skimage import io, measure33logging.getLogger('models.seg_post_model.models').setLevel(logging.ERROR)34 35SCALE = 136 37 38 39class SegmentationModule(pl.LightningModule):40 def __init__(self, use_box=True):41 super().__init__()42 self.use_box = use_box43 self.config = RunConfig() # config for stable diffusion44 self.initialize_model()45 46 47 def initialize_model(self):48 49 # load loca model50 self.loca_model = build_loca_model()51 self.loca_model.eval()52 53 self.counting_adapter = Counting(scale_factor=SCALE)54 55 ### load stable diffusion and its controller56 self.stable = load_stable_diffusion_model(config=self.config)57 self.noise_scheduler = self.stable.scheduler58 self.controller = AttentionStore(max_size=64)59 attn_utils.register_attention_control(self.stable, self.controller)60 attn_utils.register_hier_output(self.stable)61 62 ##### initialize token_emb #####63 placeholder_token = "<task-prompt>"64 self.task_token = "repetitive objects"65 # Add the placeholder token in tokenizer66 num_added_tokens = self.stable.tokenizer.add_tokens(placeholder_token)67 if num_added_tokens == 0:68 raise ValueError(69 f"The tokenizer already contains the token {placeholder_token}. Please pass a different"70 " `placeholder_token` that is not already in the tokenizer."71 )72 try:73 task_embed_from_pretrain = hf_hub_download(74 repo_id="phoebe777777/111",75 filename="task_embed.pth",76 token=None,77 force_download=False78 )79 placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(placeholder_token)80 self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))81 82 token_embeds = self.stable.text_encoder.get_input_embeddings().weight.data83 token_embeds[placeholder_token_id] = task_embed_from_pretrain84 except:85 initializer_token = "segment"86 token_ids = self.stable.tokenizer.encode(initializer_token, add_special_tokens=False)87 # Check if initializer_token is a single token or a sequence of tokens88 if len(token_ids) > 1:89 # raise ValueError("The initializer token must be a single token.")90 token_ids = token_ids[:1]91 92 initializer_token_id = token_ids[0]93 placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(placeholder_token)94 95 self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))96 97 token_embeds = self.stable.text_encoder.get_input_embeddings().weight.data98 token_embeds[placeholder_token_id] = token_embeds[initializer_token_id]99 100 # others101 self.placeholder_token = placeholder_token102 self.placeholder_token_id = placeholder_token_id103 104 105 106 107 def move_to_device(self, device):108 self.stable.to(device)109 self.counting_adapter.to(device)110 self.loca_model.to(device)111 112 self.to(device)113 114 115 def forward(self, data_path, box=None):116 filename = data_path.split("/")[-1]117 img = Image.open(data_path).convert("RGB")118 width, height = img.size119 input_image = T.Compose([T.ToTensor(), T.Resize((512, 512))])(img)120 input_image_stable = input_image - 0.5121 input_image = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(input_image)122 if box is not None:123 boxes = torch.tensor(box) / torch.tensor([width, height, width, height]) * 512 # xyxy, normalized124 assert self.use_box == True125 else:126 boxes = torch.tensor([[0,0,512,512]])127 assert self.use_box == False128 img_raw = io.imread(data_path)129 if len(img_raw.shape) == 3 and img_raw.shape[2] > 3:130 img_raw = img_raw[:,:,:3]131 img_raw = cv2.resize(img_raw, (512, 512))132 133 # move to device134 input_image = input_image.unsqueeze(0).to(self.device)135 img_raw = torch.from_numpy(img_raw).unsqueeze(0).float().to(self.device)136 boxes = boxes.unsqueeze(0).to(self.device)137 input_image_stable = input_image_stable.unsqueeze(0).to(self.device)138 139 latents = self.stable.vae.encode(input_image_stable).latent_dist.sample().detach()140 latents = latents * 0.18215141 # Sample noise that we'll add to the latents142 noise = torch.randn_like(latents)143 bsz = latents.shape[0]144 timesteps = torch.tensor([20], device=latents.device).long()145 noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)146 input_ids_ = self.stable.tokenizer(147 self.placeholder_token + " " + self.task_token,148 padding="max_length",149 truncation=True,150 max_length=self.stable.tokenizer.model_max_length,151 return_tensors="pt",152 )153 input_ids = input_ids_["input_ids"].to(self.device)154 attention_mask = input_ids_["attention_mask"].to(self.device)155 encoder_hidden_states = self.stable.text_encoder(input_ids, attention_mask)[0]156 encoder_hidden_states = encoder_hidden_states.repeat(bsz, 1, 1)157 158 159 160 task_loc_idx = torch.nonzero(input_ids == self.placeholder_token_id)161 162 if self.use_box:163 loca_out = self.loca_model.forward_before_reg(input_image, boxes)164 loca_feature_bf_regression = loca_out["feature_bf_regression"]165 adapted_emb = self.counting_adapter.adapter(loca_feature_bf_regression, boxes) # shape [1, 768]166 # adapted_emb = self.counting_adapter.adapter(data['crops_dino'], self.dino) # shape [1, 768]167 if task_loc_idx.shape[0] == 0:168 encoder_hidden_states[0,2,:] = adapted_emb.squeeze() # 放在task prompt下一位169 else:170 encoder_hidden_states[:,task_loc_idx[0, 1]+1,:] = adapted_emb.squeeze() # 放在task prompt下一位171 172 # Predict the noise residual173 noise_pred, feature_list = self.stable.unet(noisy_latents, timesteps, encoder_hidden_states)174 time3 = time.time()175 noise_pred = noise_pred.sample176 177 attention_store = self.controller.attention_store178 179 180 attention_maps = []181 exemplar_attention_maps1 = []182 exemplar_attention_maps2 = []183 exemplar_attention_maps3 = []184 185 # only use 64x64 self-attention186 self_attn_aggregate = attn_utils.aggregate_attention( # [res, res, 4096]187 prompts=[self.config.prompt for i in range(bsz)], # 这里要改么188 attention_store=self.controller, 189 res=64,190 from_where=("up", "down"),191 is_cross=False,192 select=0193 )194 195 # cross attention196 for res in [32, 16]:197 attn_aggregate = attn_utils.aggregate_attention( # [res, res, 77]198 prompts=[self.config.prompt for i in range(bsz)], # 这里要改么199 attention_store=self.controller, 200 res=res,201 from_where=("up", "down"),202 is_cross=True,203 select=0204 )205 206 task_attn_ = attn_aggregate[:, :, 1].unsqueeze(0).unsqueeze(0) # [1, 1, res, res]207 attention_maps.append(task_attn_)208 exemplar_attns1 = attn_aggregate[:, :, 2].unsqueeze(0).unsqueeze(0) # 取exemplar的attn209 exemplar_attention_maps1.append(exemplar_attns1)210 exemplar_attns2 = attn_aggregate[:, :, 3].unsqueeze(0).unsqueeze(0) # 取exemplar的attn211 exemplar_attention_maps2.append(exemplar_attns2)212 exemplar_attns3 = attn_aggregate[:, :, 4].unsqueeze(0).unsqueeze(0) # 取exemplar的attn213 exemplar_attention_maps3.append(exemplar_attns3)214 215 216 scale_factors = [(64 // attention_maps[i].shape[-1]) for i in range(len(attention_maps))]217 attns = torch.cat([F.interpolate(attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(attention_maps))])218 task_attn_64 = torch.mean(attns, dim=0, keepdim=True)219 cross_self_task_attn = attn_utils.self_cross_attn(self_attn_aggregate, task_attn_64)220 task_attn_64 = (task_attn_64 - task_attn_64.min()) / (task_attn_64.max() - task_attn_64.min() + 1e-6)221 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)222 223 scale_factors = [(64 // exemplar_attention_maps1[i].shape[-1]) for i in range(len(exemplar_attention_maps1))]224 attns = torch.cat([F.interpolate(exemplar_attention_maps1[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps1))])225 exemplar_attn_64_1 = torch.mean(attns, dim=0, keepdim=True)226 227 if self.use_box:228 exemplar_attn_64 = exemplar_attn_64_1229 cross_self_exe_attn = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64)230 exemplar_attn_64 = (exemplar_attn_64 - exemplar_attn_64.min()) / (exemplar_attn_64.max() - exemplar_attn_64.min() + 1e-6)231 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)232 else:233 234 scale_factors = [(64 // exemplar_attention_maps2[i].shape[-1]) for i in range(len(exemplar_attention_maps2))]235 attns = torch.cat([F.interpolate(exemplar_attention_maps2[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps2))])236 exemplar_attn_64_2 = torch.mean(attns, dim=0, keepdim=True)237 238 scale_factors = [(64 // exemplar_attention_maps3[i].shape[-1]) for i in range(len(exemplar_attention_maps3))]239 attns = torch.cat([F.interpolate(exemplar_attention_maps3[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps3))])240 exemplar_attn_64_3 = torch.mean(attns, dim=0, keepdim=True)241 242 cross_self_exe_attn1 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_1)243 cross_self_exe_attn2 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_2)244 cross_self_exe_attn3 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_3)245 # # average246 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)247 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)248 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)249 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)250 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)251 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)252 253 exemplar_attn_64 = (exemplar_attn_64_1 + exemplar_attn_64_2 + exemplar_attn_64_3) / 3254 cross_self_exe_attn = (cross_self_exe_attn1 + cross_self_exe_attn2 + cross_self_exe_attn3) / 3255 256 257 258 259 260 if self.use_box:261 attn_stack = [task_attn_64 / 2, cross_self_task_attn / 2, exemplar_attn_64, cross_self_exe_attn]262 else:263 attn_stack = [exemplar_attn_64 / 2, cross_self_exe_attn / 2, exemplar_attn_64, cross_self_exe_attn]264 attn_stack = torch.cat(attn_stack, dim=1)265 266 267 attn_after_new_regressor = self.counting_adapter.regressor(img_raw, attn_stack, feature_list) # 直接用自己的268 269 input_image = cv2.resize(input_image[0].permute(1,2,0).cpu().numpy(), (width, height))270 pred = cv2.resize(attn_after_new_regressor.squeeze().cpu().numpy(), (width, height), interpolation=cv2.INTER_NEAREST)271 return pred272 273 274 275 276 277def inference(data_path, box=None, save_path="./example_imgs", visualize=False):278 if box is not None:279 use_box = True280 else:281 use_box = False282 model = SegmentationModule(use_box=use_box)283 load_msg = model.load_state_dict(torch.load("pretrained/microscopy_matching_seg.pth"), strict=True)284 model.eval()285 with torch.no_grad():286 mask = model(data_path, box)287 288 289 # visualize290 if visualize:291 img = io.imread(data_path)292 if len(img.shape) == 3 and img.shape[2] > 3:293 img = img[:,:,:3]294 if len(img.shape) == 2:295 img = np.stack([img]*3, axis=-1)296 img_show = img.squeeze()297 mask_show = mask.squeeze()298 os.makedirs(save_path, exist_ok=True)299 filename = data_path.split("/")[-1]300 fig, ax = plt.subplots(1,2, figsize=(12,6))301 ax[0].imshow(img_show)302 if use_box:303 boxes = np.array(box)304 for box in boxes:305 rect = patches.Rectangle((box[0], box[1]), box[2]-box[0], box[3]-box[1], linewidth=2, edgecolor='r', facecolor='none')306 ax[0].add_patch(rect)307 ax[0].set_title("Input Image with Box")308 else:309 ax[0].set_title("Input Image")310 ax[0].axis("off")311 ax[1].imshow(img_show)312 for inst_id in np.unique(mask_show):313 if inst_id == 0: # 0 background314 continue315 # 生成二值 mask316 binary_mask = (mask_show == inst_id).astype(np.uint8)317 contours = measure.find_contours(binary_mask, 0.5)318 for contour in contours:319 ax[1].plot(contour[:, 1], contour[:, 0], linewidth=1.5, linestyle="--", color='yellow')320 ax[1].imshow(overlay_instances(img_show, mask_show, alpha=0.3))321 ax[1].set_title("Segmentation Result")322 ax[1].axis("off")323 plt.tight_layout()324 plt.savefig(os.path.join(save_path, filename.split(".")[0]+"_seg.png"), dpi=300)325 plt.close()326 327 return mask328 329 330def main():331 inference(332 data_path="example_imgs/1977_Well_F-5_Field_1.png", 333 # box=[[724, 864, 900, 966]], 334 save_path="./example_imgs",335 visualize=True336 )337 338 339from matplotlib import cm340 341def overlay_instances(img, mask, alpha=0.5, cmap_name="tab20"):342 img = img.astype(np.float32)343 if len(img.shape) == 2:344 img = np.stack([img]*3, axis=-1)345 if img.max() > 1.5:346 img = img / 255.0347 348 349 overlay = img.copy()350 cmap = cm.get_cmap(cmap_name, np.max(mask)+1)351 352 for inst_id in np.unique(mask):353 if inst_id == 0: 354 continue355 color = np.array(cmap(inst_id)[:3]) # RGB356 overlay[mask == inst_id] = (1 - alpha) * overlay[mask == inst_id] + alpha * color357 358 return overlay359 360if __name__ == "__main__":361 main()