VisionLanguageGroup/MicroscopyMatching

1# stable diffusion x loca2import os3import pprint4from typing import Any, List, Optional5import argparse6from huggingface_hub import hf_hub_download7import pyrallis8from pytorch_lightning.utilities.types import STEP_OUTPUT9import torch10import os11from PIL import Image12import numpy as np13from config import RunConfig14from _utils import attn_utils_new as attn_utils15from _utils.attn_utils_new import AttentionStore16from _utils.misc_helper import *17import torch.nn.functional as F18import matplotlib.pyplot as plt19import cv220import warnings21warnings.filterwarnings("ignore", category=UserWarning)22import pytorch_lightning as pl23from _utils.load_models import load_stable_diffusion_model24from models.model import Counting_with_SD_features_loca as Counting25from models.enc_model.loca import build_model as build_loca_model26import time27import torchvision.transforms as T28import skimage.io as io29 30SCALE = 131 32 33class CountingModule(pl.LightningModule):34    def __init__(self, use_box=True):35        super().__init__()36        self.use_box = use_box37        self.config = RunConfig()   # config for stable diffusion38        self.initialize_model()39        40 41    def initialize_model(self):42        43        # load loca model44        self.loca_model = build_loca_model()45 46        self.counting_adapter = Counting(scale_factor=SCALE)47        # if os.path.isfile(self.args.adapter_weight):48        #     adapter_weight = torch.load(self.args.adapter_weight,map_location=torch.device('cpu'))49        #     self.counting_adapter.load_state_dict(adapter_weight, strict=False)50            51        ### load stable diffusion and its controller52        self.stable = load_stable_diffusion_model(config=self.config)53        self.noise_scheduler = self.stable.scheduler54        self.controller = AttentionStore(max_size=64)55        attn_utils.register_attention_control(self.stable, self.controller)56        attn_utils.register_hier_output(self.stable)57 58        ##### initialize token_emb #####59        placeholder_token = "<task-prompt>"60        self.task_token = "repetitive objects"61        # Add the placeholder token in tokenizer62        num_added_tokens = self.stable.tokenizer.add_tokens(placeholder_token)63        if num_added_tokens == 0:64            raise ValueError(65                f"The tokenizer already contains the token {placeholder_token}. Please pass a different"66                " `placeholder_token` that is not already in the tokenizer."67            )68        try:69            task_embed_from_pretrain = hf_hub_download(70                repo_id="phoebe777777/111",71                filename="task_embed.pth",72                token=None,73                force_download=False74            )75            placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(placeholder_token)76            self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))77 78            token_embeds = self.stable.text_encoder.get_input_embeddings().weight.data79            token_embeds[placeholder_token_id] = task_embed_from_pretrain80        except:81            initializer_token = "count"82            token_ids = self.stable.tokenizer.encode(initializer_token, add_special_tokens=False)83            # Check if initializer_token is a single token or a sequence of tokens84            if len(token_ids) > 1:85                raise ValueError("The initializer token must be a single token.")86 87            initializer_token_id = token_ids[0]88            placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(placeholder_token)89 90            self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))91 92            token_embeds = self.stable.text_encoder.get_input_embeddings().weight.data93            token_embeds[placeholder_token_id] = token_embeds[initializer_token_id]94 95        # others96        self.placeholder_token = placeholder_token97        self.placeholder_token_id = placeholder_token_id98    99 100    def move_to_device(self, device):101        self.stable.to(device)102        if self.loca_model is not None and self.counting_adapter is not None:103            self.loca_model.to(device)104            self.counting_adapter.to(device)105        self.to(device)106 107    def forward(self, data_path, box=None):108        filename = data_path.split("/")[-1]109        img = Image.open(data_path).convert("RGB")110        width, height = img.size111        input_image = T.Compose([T.ToTensor(), T.Resize((512, 512))])(img)112        input_image_stable = input_image - 0.5113        input_image = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(input_image)114        if box is not None:115            boxes = torch.tensor(box) / torch.tensor([width, height, width, height]) * 512  # xyxy, normalized116            assert self.use_box == True117        else:118            boxes = torch.tensor([[100,100,130,130], [200,200,250,250]], dtype=torch.float32)  # dummy box119            assert self.use_box == False120 121        # move to device122        input_image = input_image.unsqueeze(0).to(self.device)123        boxes = boxes.unsqueeze(0).to(self.device)124        input_image_stable = input_image_stable.unsqueeze(0).to(self.device)125        126 127 128        latents = self.stable.vae.encode(input_image_stable).latent_dist.sample().detach()129        latents = latents * 0.18215130        # Sample noise that we'll add to the latents131        noise = torch.randn_like(latents)132        timesteps = torch.tensor([20], device=latents.device).long()133        noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)134        input_ids_ = self.stable.tokenizer(135            self.placeholder_token + " repetitive objects",136            # "object",137            padding="max_length",138            truncation=True,139            max_length=self.stable.tokenizer.model_max_length,140            return_tensors="pt",141        )142        input_ids = input_ids_["input_ids"].to(self.device)143        attention_mask = input_ids_["attention_mask"].to(self.device)144        encoder_hidden_states = self.stable.text_encoder(input_ids, attention_mask)[0]145 146        input_image = input_image.to(self.device)147        boxes = boxes.to(self.device)148        149 150        task_loc_idx = torch.nonzero(input_ids == self.placeholder_token_id)151        if self.use_box:152            loca_out = self.loca_model.forward_before_reg(input_image, boxes)153            loca_feature_bf_regression =  loca_out["feature_bf_regression"]154            adapted_emb = self.counting_adapter.adapter(loca_feature_bf_regression, boxes)      # shape [1, 768]155            if task_loc_idx.shape[0] == 0:156                encoder_hidden_states[0,2,:] = adapted_emb.squeeze()  157            else:158                encoder_hidden_states[0,task_loc_idx[0, 1]+1,:] = adapted_emb.squeeze()  159 160        # Predict the noise residual161        noise_pred, feature_list = self.stable.unet(noisy_latents, timesteps, encoder_hidden_states)162        noise_pred = noise_pred.sample163        attention_store = self.controller.attention_store164 165 166        attention_maps = []167        exemplar_attention_maps = []168        exemplar_attention_maps1 = []169        exemplar_attention_maps2 = []170        exemplar_attention_maps3 = []171 172        cross_self_task_attn_maps = []173        cross_self_exe_attn_maps = []174 175        # only use 64x64 self-attention176        self_attn_aggregate = attn_utils.aggregate_attention( # [res, res, 4096]177                prompts=[self.config.prompt],        178                attention_store=self.controller,     179                res=64,180                from_where=("up", "down"),181                is_cross=False,182                select=0183            )184        self_attn_aggregate32 = attn_utils.aggregate_attention( # [res, res, 4096]185                prompts=[self.config.prompt],        186                attention_store=self.controller,     187                res=32,188                from_where=("up", "down"),189                is_cross=False,190                select=0191            )192        self_attn_aggregate16 = attn_utils.aggregate_attention( # [res, res, 4096]193                prompts=[self.config.prompt],        194                attention_store=self.controller,     195                res=16,196                from_where=("up", "down"),197                is_cross=False,198                select=0199            )200 201        # cross attention202        for res in [32, 16]:203            attn_aggregate = attn_utils.aggregate_attention( # [res, res, 77]204                prompts=[self.config.prompt],        205                attention_store=self.controller,     206                res=res,207                from_where=("up", "down"),208                is_cross=True,209                select=0210            )211 212            task_attn_ = attn_aggregate[:, :, 1].unsqueeze(0).unsqueeze(0) # [1, 1, res, res]213            attention_maps.append(task_attn_)214            if self.use_box:215                exemplar_attns = attn_aggregate[:, :, 2].unsqueeze(0).unsqueeze(0) 216                exemplar_attention_maps.append(exemplar_attns)217            else:218                exemplar_attns1 = attn_aggregate[:, :, 2].unsqueeze(0).unsqueeze(0)219                exemplar_attns2 = attn_aggregate[:, :, 3].unsqueeze(0).unsqueeze(0)220                exemplar_attns3 = attn_aggregate[:, :, 4].unsqueeze(0).unsqueeze(0)221                exemplar_attention_maps1.append(exemplar_attns1)222                exemplar_attention_maps2.append(exemplar_attns2)223                exemplar_attention_maps3.append(exemplar_attns3)224 225 226        scale_factors = [(64 // attention_maps[i].shape[-1]) for i in range(len(attention_maps))]227        attns = torch.cat([F.interpolate(attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(attention_maps))])228        task_attn_64 = torch.mean(attns, dim=0, keepdim=True)229        cross_self_task_attn = attn_utils.self_cross_attn(self_attn_aggregate, task_attn_64)230        cross_self_task_attn_maps.append(cross_self_task_attn)231 232        if self.use_box:233            scale_factors = [(64 // exemplar_attention_maps[i].shape[-1]) for i in range(len(exemplar_attention_maps))]234            attns = torch.cat([F.interpolate(exemplar_attention_maps[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps))])235            exemplar_attn_64 = torch.mean(attns, dim=0, keepdim=True)236 237            cross_self_exe_attn = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64)238            cross_self_exe_attn_maps.append(cross_self_exe_attn)239        else:240            scale_factors = [(64 // exemplar_attention_maps1[i].shape[-1]) for i in range(len(exemplar_attention_maps1))]241            attns = torch.cat([F.interpolate(exemplar_attention_maps1[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps1))])242            exemplar_attn_64_1 = torch.mean(attns, dim=0, keepdim=True)243 244            scale_factors = [(64 // exemplar_attention_maps2[i].shape[-1]) for i in range(len(exemplar_attention_maps2))]245            attns = torch.cat([F.interpolate(exemplar_attention_maps2[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps2))])246            exemplar_attn_64_2 = torch.mean(attns, dim=0, keepdim=True)247 248            scale_factors = [(64 // exemplar_attention_maps3[i].shape[-1]) for i in range(len(exemplar_attention_maps3))]249            attns = torch.cat([F.interpolate(exemplar_attention_maps3[i_], scale_factor=scale_factors[i_], mode="bilinear") for i_ in range(len(exemplar_attention_maps3))])250            exemplar_attn_64_3 = torch.mean(attns, dim=0, keepdim=True)251 252 253            cross_self_task_attn = attn_utils.self_cross_attn(self_attn_aggregate, task_attn_64)254            cross_self_task_attn_maps.append(cross_self_task_attn)255 256            # if self.args.merge_exemplar == "average":257            cross_self_exe_attn1 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_1)258            cross_self_exe_attn2 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_2)259            cross_self_exe_attn3 = attn_utils.self_cross_attn(self_attn_aggregate, exemplar_attn_64_3)260            exemplar_attn_64 = (exemplar_attn_64_1 + exemplar_attn_64_2 + exemplar_attn_64_3) / 3261            cross_self_exe_attn = (cross_self_exe_attn1 + cross_self_exe_attn2 + cross_self_exe_attn3) / 3262 263            exemplar_attn_64 = (exemplar_attn_64 - exemplar_attn_64.min()) / (exemplar_attn_64.max() - exemplar_attn_64.min() + 1e-6)264 265        attn_stack = [exemplar_attn_64 / 2, cross_self_exe_attn / 2, exemplar_attn_64, cross_self_exe_attn]266        attn_stack = torch.cat(attn_stack, dim=1)267        268        if not self.use_box:269 270            loca_out = self.loca_model.forward_before_reg(input_image, boxes)271            loca_feature_bf_regression =  loca_out["feature_bf_regression"]272        attn_out = self.loca_model.forward_reg(loca_out, attn_stack, feature_list[-1])273        pred_density = attn_out["pred"].squeeze().cpu().numpy()274        pred_cnt = pred_density.sum().item()275 276        # resize pred_density to original image size277        pred_density_rsz = cv2.resize(pred_density, (width, height), interpolation=cv2.INTER_CUBIC)278        pred_density_rsz = pred_density_rsz / pred_density_rsz.sum() * pred_cnt279 280        return pred_density_rsz, pred_cnt281 282    283def inference(data_path, box=None, save_path="./example_imgs", visualize=False):284    if box is not None:285        use_box = True286    else:287        use_box = False288    model = CountingModule(use_box=use_box)289    load_msg = model.load_state_dict(torch.load("pretrained/microscopy_matching_cnt.pth"), strict=True)290    model.eval()291    with torch.no_grad():292        density_map, cnt = model(data_path, box)293    294    if visualize:295        img = io.imread(data_path)296        if len(img.shape) == 3 and img.shape[2] > 3:297            img = img[:,:,:3]298        if len(img.shape) == 2:299            img = np.stack([img]*3, axis=-1)300        img_show = img.squeeze()301        density_map_show = density_map.squeeze()302        os.makedirs(save_path, exist_ok=True)303        filename = data_path.split("/")[-1]304        img_show = (img_show - np.min(img_show)) / (np.max(img_show) - np.min(img_show))305        fig, ax = plt.subplots(1,2, figsize=(12,6))306        ax[0].imshow(img_show)307        ax[0].axis('off')308        ax[0].set_title(f"Input image")309        ax[1].imshow(img_show)310        ax[1].imshow(density_map_show, cmap='jet', alpha=0.5)  # Overlay density map with some transparency311        ax[1].axis('off')312        ax[1].set_title(f"Predicted density map, count: {cnt:.1f}")313        plt.tight_layout()314        plt.savefig(os.path.join(save_path, filename.split(".")[0]+"_cnt.png"), dpi=300)315        plt.close()316    return density_map317 318def main():319    320    inference(321        data_path = "example_imgs/1977_Well_F-5_Field_1.png",322        # box=[[150, 60, 183, 87]],323        save_path = "./example_imgs",324        visualize = True325    )326 327if __name__ == "__main__":328    main()