from segment_anything import build_sam_vit_h,build_sam_vit_b
from segment_anything.utils.transforms import ResizeLongestSide
import cv2

import torch.nn.functional as F
import torch.nn as nn
import torch
pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
def preprocess_pad(x):
    x = (x - pixel_mean) / pixel_std
    # Pad
    h, w = x.shape[-2:]
    padh = 1024 - h
    padw = 1024 - w
    x = F.pad(x, (0, padw, 0, padh))
    return x
transform = ResizeLongestSide(1024)
img_path = './data/reason_seg/train/16438974836_22b8cca346_o.jpg'
sam_image = transform.apply_image(cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB))
sam_image = preprocess_pad(torch.from_numpy(sam_image).permute(2, 0, 1).contiguous())
def initialize_grounding_modules():
        # SAM
        # self.visual_model = build_sam_vit_h(self.vision_pretrained)
        # import ipdb;ipdb.set_trace()
        # visual_model = build_sam_vit_b('./checkpoints/sam_vit_b_01ec64.pth').to(device='cuda:0', dtype=torch.bfloat16)
        visual_model = build_sam_vit_b('./checkpoints/SAM.pth').to(device='cuda:0', dtype=torch.bfloat16)
        
        for param in visual_model.parameters():
            param.requires_grad = False
        train_mask_decoder=True 
        if train_mask_decoder:
            visual_model.mask_decoder.train()
            for param in visual_model.mask_decoder.parameters():
                param.requires_grad = True
        # Projection layer
        in_dim = 2048
        out_dim = 256
        text_fc = [
            nn.Linear(in_dim, in_dim),
            nn.ReLU(inplace=True),
            nn.Linear(in_dim, out_dim),
            nn.Dropout(0.0),
        ]
        text_hidden_fcs = nn.ModuleList([nn.Sequential(*text_fc)])
        text_hidden_fcs.train()
        for param in text_hidden_fcs.parameters():
            param.requires_grad = True
        return visual_model
visual_model=initialize_grounding_modules()
def get_visual_embs( pixel_values: torch.FloatTensor):
    with torch.no_grad():
        image_embeddings_list = []
        for i in range(pixel_values.shape[0]):
            torch.cuda.empty_cache()
            image_embeddings = visual_model.image_encoder(
                pixel_values[i].unsqueeze(0)
            )
            image_embeddings_list.append(image_embeddings)
        torch.cuda.empty_cache()
        image_embeddings = torch.cat(image_embeddings_list, 0)
    return image_embeddings
get_visual_embs(sam_image.unsqueeze(0).to(device='cuda:0', dtype=torch.bfloat16)) #[1, 256, 64, 64]
import matplotlib.pyplot as plt
tmp = sam_image.cpu().float()
image_np = tmp.numpy()
image_np = image_np.transpose(1, 2, 0)
min_val = image_np.min()
max_val = image_np.max()
image_normalized = (image_np - min_val) / (max_val - min_val)
plt.figure(figsize=(8, 8))  # 设置图像大小
plt.imshow(image_normalized)  # 显示归一化后的图像
plt.axis('off')  # 隐藏坐标轴
plt.tight_layout()  # 去除周围空白
plt.savefig('visualized_image_sam.png', bbox_inches='tight', dpi=300)