Upload 29 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+GOAL_github/datasets/docci_segment_sim_bbox_del_org.json filter=lfs diff=lfs merge=lfs -text

GOAL_github/LISM/DCIsamwithclipmaxsimmaintainbackrefsentencefortestset.py

@@ -0,0 +1,170 @@
+import requests
+from io import BytesIO
+from PIL import Image
+import numpy as np
+from transformers import pipeline, CLIPProcessor, CLIPModel
+import cv2
+import torch
+import json
+import os
+from tqdm import tqdm
+
+# SAM 모델 로드
+generator = pipeline("mask-generation", device=2, points_per_batch=128)
+
+# CLIP 모델 로드
+model_name = "openai/clip-vit-large-patch14-336"
+clip_model = CLIPModel.from_pretrained(model_name)
+clip_processor = CLIPProcessor.from_pretrained(model_name)
+
+# 원본 배경 유지
+def keep_original_background(image, mask):
+    masked_image = np.array(image)
+    mask_3d = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
+    masked_image = masked_image * mask_3d + (1 - mask_3d) * np.array(image)
+    return Image.fromarray(np.uint8(masked_image))
+
+# 이미지 리사이징 함수
+def resize_image_if_needed(image_path, max_size=(2048, 1536)):
+    with Image.open(image_path) as img:
+        if img.size[0] > 4000 or img.size[1] > 3000:
+            img.thumbnail(max_size, Image.LANCZOS)
+            return img.copy()
+        return img.copy()
+
+# CLIP을 사용한 이미지-캡션 매칭
+def get_clip_similarity(images, texts):
+    try:
+        inputs = clip_processor(text=texts, images=images, return_tensors="pt", padding=True)
+        outputs = clip_model(**inputs)
+        return outputs.logits_per_image.cpu().detach().numpy()
+    except Exception as e:
+        print(f"Error in CLIP processing: {str(e)}")
+        return None
+
+# 간단한 문장 토큰화
+def tokenize_caption(text):
+    return [s.strip() for s in text.split('.') if s.strip()]
+
+# 기본 경로 설정
+matching_results_path = "matching_results"
+output_base_dir = "segment_with_background_DCI_test_set_max_0.01"
+os.makedirs(output_base_dir, exist_ok=True)
+
+# train_sa로 시작하는 폴더들 찾기
+train_folders = [f for f in os.listdir(matching_results_path) if f.startswith('test_sa_')]
+
+# 각 train 폴더에 대해 처리
+for train_folder in tqdm(train_folders, desc="Processing folders"):
+    folder_path = os.path.join(matching_results_path, train_folder)
+    
+    # JSON 파일 찾기 (sa_xxxxxx_result.json)
+    json_files = [f for f in os.listdir(folder_path) if f.endswith('_result.json')]
+    if not json_files:
+        continue
+    
+    json_path = os.path.join(folder_path, json_files[0])
+    
+    # JSON 파일 로드
+    try:
+        with open(json_path, 'r') as f:
+            annotation = json.load(f)
+    except Exception as e:
+        print(f"Error loading JSON for {train_folder}: {str(e)}")
+        continue
+
+    image_filename = annotation['original_image']
+    image_path = os.path.join(folder_path, image_filename)
+    caption = annotation['extra_caption']
+
+    # 결과를 저장할 디렉토리 설정
+    output_dir = os.path.join(output_base_dir, f"{image_filename}_results")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 이미 처리된 이미지는 건너뛰기
+    if os.path.exists(os.path.join(output_dir, 'matched_dataset_test.json')):
+        continue
+
+    # 이미지 로드 및 세그멘테이션
+    try:
+        original_image = Image.open(image_path)
+        resized_image = resize_image_if_needed(image_path)
+        
+        # 리사이즈된 이미지로 SAM 실행
+        outputs = generator(resized_image, points_per_batch=128)
+        
+        # 원본 이미지 크기로 마스크 리사이즈
+        if original_image.size != resized_image.size:
+            for i in range(len(outputs['masks'])):
+                mask = Image.fromarray(outputs['masks'][i])
+                mask = mask.resize(original_image.size, Image.NEAREST)
+                outputs['masks'][i] = np.array(mask)
+        
+        image = original_image
+    except Exception as e:
+        print(f"Error processing {image_filename}: {str(e)}")
+        continue
+
+    # 필터링 및 세그먼트 이미지 생성
+    min_area_ratio = 0.01
+    max_area_ratio = 0.8
+    total_area = image.size[0] * image.size[1]
+    segmented_images = [image]
+    segmented_image_paths = ["original_image.jpg"]
+
+    for i, mask in enumerate(outputs['masks']):
+        segment_area = np.sum(mask)
+        area_ratio = segment_area / total_area
+        if min_area_ratio <= area_ratio <= max_area_ratio:
+            masked_image = keep_original_background(image, mask)
+            
+            y, x = np.where(mask)
+            y_min, y_max, x_min, x_max = y.min(), y.max(), x.min(), x.max()
+            cropped_image = np.array(masked_image)[y_min:y_max+1, x_min:x_max+1]
+            
+            segmented_images.append(Image.fromarray(np.uint8(cropped_image)))
+            segmented_image_paths.append(f"{image_filename.split('.')[0]}_max_{i+1}.png")
+
+    # 캡션 토큰화
+    tokenized_captions = tokenize_caption(caption)
+
+    # 모든 문장-이미지 쌍에 대한 유사도 계산
+    similarity_matrix = get_clip_similarity(segmented_images, tokenized_captions)
+
+    if similarity_matrix is None:
+        print(f"Skipping image {image_filename} due to CLIP processing error")
+        continue
+
+    matched_data = []
+    saved_images = set()
+
+    # 모든 문장에 대해 가장 높은 유사도를 가진 이미지와 매칭
+    for j, caption in enumerate(tokenized_captions):
+        i = np.argmax(similarity_matrix[:, j])
+        similarity_score = float(similarity_matrix[i, j])
+        matched_image_path = segmented_image_paths[i]
+        
+        matched_data.append({
+            "caption": caption,
+            "matched_image_path": matched_image_path,
+            "similarity_score": similarity_score
+        })
+        
+        saved_images.add(i)
+
+    # 원본 이미지 저장
+    original_image_path = os.path.join(output_dir, "original_image.jpg")
+    image.save(original_image_path)
+
+    # 매칭된 세그먼트 이미지 저장
+    for i in saved_images:
+        if i != 0:  # 원본 이미지는 건너뛰기
+            output_file_path = os.path.join(output_dir, segmented_image_paths[i])
+            segmented_images[i].save(output_file_path)
+
+    # 결과 저장
+    json_output_path = os.path.join(output_dir, 'matched_dataset_test.json')
+    with open(json_output_path, 'w') as f:
+        json.dump(matched_data, f, indent=2)
+
+print("모든 test set 이미지 처리가 완료되었습니다.")

GOAL_github/LISM/DCIsamwithclipmaxsimmaintainbackrefsentencefortrainset.py

@@ -0,0 +1,170 @@
+import requests
+from io import BytesIO
+from PIL import Image
+import numpy as np
+from transformers import pipeline, CLIPProcessor, CLIPModel
+import cv2
+import torch
+import json
+import os
+from tqdm import tqdm
+
+# SAM 모델 로드
+generator = pipeline("mask-generation", device=1, points_per_batch=128)
+
+# CLIP 모델 로드
+model_name = "openai/clip-vit-large-patch14-336"
+clip_model = CLIPModel.from_pretrained(model_name)
+clip_processor = CLIPProcessor.from_pretrained(model_name)
+
+# 원본 배경 유지
+def keep_original_background(image, mask):
+    masked_image = np.array(image)
+    mask_3d = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
+    masked_image = masked_image * mask_3d + (1 - mask_3d) * np.array(image)
+    return Image.fromarray(np.uint8(masked_image))
+
+# 이미지 리사이징 함수
+def resize_image_if_needed(image_path, max_size=(2048, 1536)):
+    with Image.open(image_path) as img:
+        if img.size[0] > 4000 or img.size[1] > 3000:
+            img.thumbnail(max_size, Image.LANCZOS)
+            return img.copy()
+        return img.copy()
+
+# CLIP을 사용한 이미지-캡션 매칭
+def get_clip_similarity(images, texts):
+    try:
+        inputs = clip_processor(text=texts, images=images, return_tensors="pt", padding=True)
+        outputs = clip_model(**inputs)
+        return outputs.logits_per_image.cpu().detach().numpy()
+    except Exception as e:
+        print(f"Error in CLIP processing: {str(e)}")
+        return None
+
+# 간단한 문장 토큰화
+def tokenize_caption(text):
+    return [s.strip() for s in text.split('.') if s.strip()]
+
+# 기본 경로 설정
+matching_results_path = "matching_results"
+output_base_dir = "segment_with_background_DCI_train_set_max_0.01"
+os.makedirs(output_base_dir, exist_ok=True)
+
+# train_sa로 시작하는 폴더들 찾기
+train_folders = [f for f in os.listdir(matching_results_path) if f.startswith('train_sa_')]
+
+# 각 train 폴더에 대해 처리
+for train_folder in tqdm(train_folders, desc="Processing folders"):
+    folder_path = os.path.join(matching_results_path, train_folder)
+    
+    # JSON 파일 찾기 (sa_xxxxxx_result.json)
+    json_files = [f for f in os.listdir(folder_path) if f.endswith('_result.json')]
+    if not json_files:
+        continue
+    
+    json_path = os.path.join(folder_path, json_files[0])
+    
+    # JSON 파일 로드
+    try:
+        with open(json_path, 'r') as f:
+            annotation = json.load(f)
+    except Exception as e:
+        print(f"Error loading JSON for {train_folder}: {str(e)}")
+        continue
+
+    image_filename = annotation['original_image']
+    image_path = os.path.join(folder_path, image_filename)
+    caption = annotation['extra_caption']
+
+    # 결과를 저장할 디렉토리 설정
+    output_dir = os.path.join(output_base_dir, f"{image_filename}_results")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 이미 처리된 이미지는 건너뛰기
+    if os.path.exists(os.path.join(output_dir, 'matched_dataset_train.json')):
+        continue
+
+    # 이미지 로드 및 세그멘테이션
+    try:
+        original_image = Image.open(image_path)
+        resized_image = resize_image_if_needed(image_path)
+        
+        # 리사이즈된 이미지로 SAM 실행
+        outputs = generator(resized_image, points_per_batch=128)
+        
+        # 원본 이미지 크기로 마스크 리사이즈
+        if original_image.size != resized_image.size:
+            for i in range(len(outputs['masks'])):
+                mask = Image.fromarray(outputs['masks'][i])
+                mask = mask.resize(original_image.size, Image.NEAREST)
+                outputs['masks'][i] = np.array(mask)
+        
+        image = original_image
+    except Exception as e:
+        print(f"Error processing {image_filename}: {str(e)}")
+        continue
+
+    # 필터링 및 세그먼트 이미지 생성
+    min_area_ratio = 0.01
+    max_area_ratio = 0.8
+    total_area = image.size[0] * image.size[1]
+    segmented_images = [image]
+    segmented_image_paths = ["original_image.jpg"]
+
+    for i, mask in enumerate(outputs['masks']):
+        segment_area = np.sum(mask)
+        area_ratio = segment_area / total_area
+        if min_area_ratio <= area_ratio <= max_area_ratio:
+            masked_image = keep_original_background(image, mask)
+            
+            y, x = np.where(mask)
+            y_min, y_max, x_min, x_max = y.min(), y.max(), x.min(), x.max()
+            cropped_image = np.array(masked_image)[y_min:y_max+1, x_min:x_max+1]
+            
+            segmented_images.append(Image.fromarray(np.uint8(cropped_image)))
+            segmented_image_paths.append(f"{image_filename.split('.')[0]}_max_{i+1}.png")
+
+    # 캡션 토큰화
+    tokenized_captions = tokenize_caption(caption)
+
+    # 모든 문장-이미지 쌍에 대한 유사도 계산
+    similarity_matrix = get_clip_similarity(segmented_images, tokenized_captions)
+
+    if similarity_matrix is None:
+        print(f"Skipping image {image_filename} due to CLIP processing error")
+        continue
+
+    matched_data = []
+    saved_images = set()
+
+    # 모든 문장에 대해 가장 높은 유사도를 가진 이미지와 매칭
+    for j, caption in enumerate(tokenized_captions):
+        i = np.argmax(similarity_matrix[:, j])
+        similarity_score = float(similarity_matrix[i, j])
+        matched_image_path = segmented_image_paths[i]
+        
+        matched_data.append({
+            "caption": caption,
+            "matched_image_path": matched_image_path,
+            "similarity_score": similarity_score
+        })
+        
+        saved_images.add(i)
+
+    # 원본 이미지 저장
+    original_image_path = os.path.join(output_dir, "original_image.jpg")
+    image.save(original_image_path)
+
+    # 매칭된 세그먼트 이미지 저장
+    for i in saved_images:
+        if i != 0:  # 원본 이미지는 건너뛰기
+            output_file_path = os.path.join(output_dir, segmented_image_paths[i])
+            segmented_images[i].save(output_file_path)
+
+    # 결과 저장
+    json_output_path = os.path.join(output_dir, 'matched_dataset_train.json')
+    with open(json_output_path, 'w') as f:
+        json.dump(matched_data, f, indent=2)
+
+print("모든 train set 이미지 처리가 완료되었습니다.")

GOAL_github/datasets/docci_segment_sim_bbox_del_org.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ce1febc6eca26cb77936fdfb2f4fc6aa8c30eb52172505e22a015d06757a7bc
+size 35738701

GOAL_github/goal.py

@@ -0,0 +1,469 @@
+import torch
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+import json
+import argparse
+from PIL import Image
+from pathlib import Path
+from torch.utils.data import Dataset
+import lightning as L
+import transformers
+import torch.nn.functional as F
+import shutil
+import time
+import numpy as np
+from utils.func import *
+from utils.transforms import *
+from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict
+import shutil
+import math
+import random
+import wandb
+
+torch.autograd.set_detect_anomaly(True)  # Enable anomaly detection
+
+def clip_loss(sim):
+    gt = torch.arange(len(sim), dtype=torch.long, device=sim.device)
+    return (torch.nn.CrossEntropyLoss()(sim, gt) + torch.nn.CrossEntropyLoss()(sim.t(), gt)) / 2.0
+
+def get_patch_tokens_from_bbox(patch_tokens, bbox, b, original_image_size, image_size=224, patch_size=16):
+    # Get original dimensions from actual image size
+    org_width, org_height = original_image_size
+    
+    # Scale coordinates to image_size
+    x1 = int(round(bbox['x1'][b].item() * image_size / org_width))
+    y1 = int(round(bbox['y1'][b].item() * image_size / org_height))
+    x2 = int(round(bbox['x2'][b].item() * image_size / org_width))
+    y2 = int(round(bbox['y2'][b].item() * image_size / org_height))
+    
+    # Ensure coordinates are within image bounds
+    x1 = max(0, min(x1, image_size-1))
+    y1 = max(0, min(y1, image_size-1))
+    x2 = max(0, min(x2, image_size))
+    y2 = max(0, min(y2, image_size))
+    
+    # Convert to patch indices (include any patch that the bbox touches)
+    patch_x1 = x1 // patch_size
+    patch_y1 = y1 // patch_size
+    patch_x2 = (x2 + patch_size - 1) // patch_size
+    patch_y2 = (y2 + patch_size - 1) // patch_size
+    
+    # Get indices of patches
+    num_patches = (image_size // patch_size)
+    indices = []
+    for i in range(patch_y1, patch_y2):
+        for j in range(patch_x1, patch_x2):
+            indices.append(i * num_patches + j + 1)
+    
+    # Extract and pool relevant patch tokens
+    relevant_tokens = patch_tokens[:, indices, :]
+    pooled_tokens = torch.mean(relevant_tokens, dim=1)
+    
+    return pooled_tokens
+
+def get_text_tokens_from_segment(text_tokens, org_text, seg_text, processor):
+    """
+    Args:
+        text_tokens: (B, L, D) tensor of text tokens - all tokens of original text
+        org_text: original text string
+        seg_text: segment text string
+        processor: CLIP processor
+    Returns:
+        pooled_tokens: (B, D) tensor of pooled text tokens from the relevant segment
+    """
+    # Text preprocessing
+    org_text = ' '.join(org_text.split()).strip()
+    seg_text = ' '.join(seg_text.split()).strip()
+
+    # Split org_text into sentences
+    sentences = org_text.split('.')
+    sentences = [s.strip() for s in sentences if s.strip()]
+
+    # Find seg_text position
+    seg_pos = org_text.find(seg_text)
+    current_pos = 0
+    sent_idx = -1
+
+    # Find position by sentence
+    for i, sent in enumerate(sentences):
+        sent = sent.strip()
+        if sent == seg_text:
+            seg_pos = current_pos
+            sent_idx = i
+            break
+        current_pos += len(sent) + 2
+
+    assert seg_pos != -1, f"Segment text not found in original text"
+
+    # Tokenize segment text
+    seg_tokens = processor(text=seg_text,
+                         return_tensors="pt",
+                         padding=False,
+                         truncation=False)
+    seg_token_length = len(seg_tokens.input_ids[0]) - 2  # Exclude CLS, EOS tokens
+
+    if sent_idx != -1:
+        # Calculate token index based on sentence position
+        text_before = '. '.join(sentences[:sent_idx]) + ('. ' if sent_idx > 0 else '')
+        tokens_before = processor(text=text_before,
+                                return_tensors="pt",
+                                padding=False,
+                                truncation=False)
+        start_idx = len(tokens_before.input_ids[0])
+    else:
+        # Calculate token index based on string position
+        text_before = org_text[:seg_pos]
+        tokens_before = processor(text=text_before,
+                                return_tensors="pt",
+                                padding=False,
+                                truncation=False)
+        start_idx = len(tokens_before.input_ids[0])
+
+    # Adjust range considering maximum token length
+    max_length = text_tokens.shape[1]  # 248
+    if start_idx >= max_length:
+        # If segment is at a position beyond max length,
+        # extract tokens from the end, securing space equal to segment length
+        end_idx = max_length - 1
+        start_idx = max(1, end_idx - seg_token_length)  # Start from after CLS token (1)
+    else:
+        # If within normal range
+        end_idx = min(start_idx + seg_token_length, max_length - 1)
+
+    # Extract tokens
+    relevant_tokens = text_tokens[:, start_idx:end_idx, :]
+
+    # Handle case when no tokens are extracted
+    if relevant_tokens.shape[1] == 0:
+        # Fallback: use tokens from the beginning
+        relevant_tokens = text_tokens[:, 1:min(1 + seg_token_length, max_length), :]
+
+    # Pool tokens
+    pooled_tokens = torch.mean(relevant_tokens, dim=1)
+
+    return pooled_tokens
+
+class DLoader(Dataset):
+    def __init__(self, data_list, processor, new_max_token):
+        self.data_list = data_list
+        self.processor = processor
+        self.new_max_token = new_max_token
+
+    def __len__(self):
+        return len(self.data_list)
+    
+    def _load_image(self, name):
+        img = Image.open(name).convert("RGB")
+        return img, img.size  # Also return original image size
+    
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+        
+        item = self.data_list[idx]
+        org_image, org_image_size = self._load_image(item["original_filename"])  # Get original image size
+        org_caption = item["original_caption"]
+        
+        # Always select the segment with the highest similarity score
+        segment = max(item["segment"], key=lambda x: x["similarity_score"])
+        
+        seg_image = self._load_image(segment["filename"])[0]
+        seg_caption = segment["caption"]
+        bbox = segment["bbox_coordinates"]
+
+        org_data = self.processor(images=org_image, text=org_caption, return_tensors="pt", 
+                                truncation=True, padding="max_length", max_length=self.new_max_token)
+        seg_data = self.processor(images=seg_image, text=seg_caption, return_tensors="pt",
+                                truncation=True, padding="max_length", max_length=self.new_max_token)
+
+        return (org_data.pixel_values[0], org_data.input_ids[0], 
+                seg_data.pixel_values[0], seg_data.input_ids[0],
+                bbox, org_caption, seg_caption, org_image_size,
+                item["original_filename"], segment["filename"])
+
+def main(args):
+    wandb.init(project="CLIP_Training_real", config=args)   
+    
+    fabric = L.Fabric(
+        accelerator="cuda", 
+        devices=args.world_size,
+        strategy="ddp",
+        precision="bf16"
+    )
+
+    fabric.launch()
+    fabric.seed_everything(1337 + fabric.global_rank)
+
+    if fabric.global_rank == 0:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    with open(args.dataset) as f:
+        train_list = json.load(f)
+    
+    with fabric.device:
+        processor = transformers.AutoProcessor.from_pretrained(args.model)
+        model = transformers.CLIPModel.from_pretrained(args.model)
+        longclip_pos_embeddings(model, args.new_max_token)
+        
+        # Load checkpoint if provided
+        if args.ckpt:
+            if fabric.global_rank == 0:
+                print(f"Loading checkpoint from {args.ckpt}")
+            checkpoint = torch.load(args.ckpt, map_location='cpu')
+            model.load_state_dict(checkpoint)
+            if fabric.global_rank == 0:
+                print("Checkpoint loaded successfully")
+        
+        print_trainable_parameters(fabric, model)
+
+    dataset_train = DLoader(train_list, processor, args.new_max_token)
+    
+    train_loader = torch.utils.data.DataLoader(
+        dataset_train, batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+        shuffle=True,
+    )
+
+    train_loader = fabric.setup_dataloaders(train_loader)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.init_lr, weight_decay=args.weight_decay)
+    model, optimizer = fabric.setup(model, optimizer)
+    
+    train(fabric, model, optimizer, train_loader, processor)
+    
+def train(fabric: L.Fabric, model: torch.nn.Module, optimizer: torch.optim.Optimizer, train_loader, processor) -> None:
+    iter = 0
+    total_iter = len(train_loader) * args.epochs
+    
+    # Define MSE Loss
+    mse_loss = torch.nn.MSELoss()
+    
+    for epoch in range(args.epochs):
+        epoch_loss = 0.0
+        epoch_loss_org = 0.0
+        epoch_loss_seg = 0.0
+        epoch_loss_patch = 0.0
+        epoch_loss_text = 0.0
+        
+        for i, samples in enumerate(train_loader):
+            # Cosine LR
+            lr = (args.init_lr - args.min_lr) * 0.5 * (1.0 + math.cos(math.pi * iter / total_iter)) + args.min_lr
+            for param_group in optimizer.param_groups:
+                param_group["lr"] = lr
+
+            org_image, org_text, seg_image, seg_text, bbox, org_caption, seg_caption, org_image_sizes, org_image_paths, seg_image_paths = samples
+            
+            # Get all embeddings including patch tokens and sequence tokens
+            outputs = model(pixel_values=torch.cat((org_image, seg_image), dim=0),
+                        input_ids=torch.cat((org_text, seg_text), dim=0),
+                        output_hidden_states=True)
+            # print(model.text_model.embeddings.position_embedding.weight.requires_grad)
+            # Get patch tokens and text tokens
+            vision_outputs = model.vision_model(torch.cat((org_image, seg_image), dim=0), output_hidden_states=True)
+            text_outputs = model.text_model(torch.cat((org_text, seg_text), dim=0), output_hidden_states=True)
+
+            # Split embeddings for org and seg
+            batch_size = org_image.shape[0]
+            org_image_embeds, seg_image_embeds = outputs.image_embeds[:batch_size], outputs.image_embeds[batch_size:]
+            org_text_embeds, seg_text_embeds = outputs.text_embeds[:batch_size], outputs.text_embeds[batch_size:]
+
+            # Get patch tokens and text tokens from the last hidden states
+            org_patch_tokens = vision_outputs.hidden_states[-1][:batch_size]  # (B, N, D)
+            org_text_tokens = text_outputs.hidden_states[-1][:batch_size]     # (B, L, D)
+            
+            # Original CLIP loss
+            eps = 1e-8
+            x_i = batch_align(fabric, F.normalize(outputs.image_embeds + eps))
+            x_t = batch_align(fabric, F.normalize(outputs.text_embeds + eps))
+            x_i_org, x_i_seg = x_i.chunk(2)
+            x_t_org, x_t_seg = x_t.chunk(2)
+            
+            # Compute original losses
+            sim_org = model.logit_scale.exp() * x_i_org @ x_t_org.t()
+            loss_org = clip_loss(sim_org)
+            sim_seg = model.logit_scale.exp() * x_i_seg @ x_t_seg.t()
+            loss_seg = clip_loss(sim_seg)
+            
+            # Compute patch-level alignment loss
+            patch_pooled = []
+            for b in range(batch_size):
+                # org_image_sizes is converted to [width_tensor, height_tensor] format
+                # Original format: (width, height) tuple
+                img_width = org_image_sizes[0][b].item()  # b-th element from width tensor
+                img_height = org_image_sizes[1][b].item()  # b-th element from height tensor
+                img_size = (img_width, img_height)
+                    
+                pooled = get_patch_tokens_from_bbox(org_patch_tokens[b:b+1], 
+                                                bbox, 
+                                                b,
+                                                img_size,  
+                                                image_size=args.image_size, 
+                                                patch_size=16)
+                patch_pooled.append(pooled)
+
+            patch_pooled = torch.cat(patch_pooled, dim=0)
+            patch_pooled = model.vision_model.post_layernorm(patch_pooled)
+            patch_pooled = model.visual_projection(patch_pooled)
+            patch_pooled = F.normalize(patch_pooled + eps, dim=-1)
+            seg_image_embeds = F.normalize(seg_image_embeds + eps, dim=-1)
+            
+            # Compute patch alignment loss with cosine similarity directly
+            sim_patch = patch_pooled @ seg_image_embeds.t()  # removed logit_scale
+            patch_diag = torch.diag(sim_patch)
+            loss_patch = mse_loss(patch_diag, torch.ones_like(patch_diag))
+            
+            # Compute text-level alignment loss
+            text_pooled = []
+            for b in range(batch_size):
+                #print(f"\nBatch {b} Text Sequences:")
+                
+                # Full token IDs of org_text
+                org_tokens = processor(text=org_caption[b], 
+                                    return_tensors="pt", 
+                                    padding=False, 
+                                    truncation=False)
+                org_token_ids = org_tokens.input_ids[0]
+                
+                # Full token IDs of seg_text
+                seg_tokens = processor(text=seg_caption[b], 
+                                    return_tensors="pt", 
+                                    padding=False, 
+                                    truncation=False)
+                seg_token_ids = seg_tokens.input_ids[0]
+                
+                # Decode token IDs to text
+                org_tokens_text = processor.tokenizer.convert_ids_to_tokens(org_token_ids)
+                seg_tokens_text = processor.tokenizer.convert_ids_to_tokens(seg_token_ids)
+                
+                # Confirm position of tokens extracted by get_text_tokens_from_segment function
+                start_idx = len(processor(text=org_caption[b][:org_caption[b].find(seg_caption[b])], 
+                                        return_tensors="pt", 
+                                        padding=False, 
+                                        truncation=False).input_ids[0])
+                
+                end_idx = start_idx + len(seg_tokens.input_ids[0]) - 2  # Exclude CLS, EOS tokens
+                
+                pooled = get_text_tokens_from_segment(org_text_tokens[b:b+1], 
+                                                    org_caption[b], 
+                                                    seg_caption[b],
+                                                    processor)
+                text_pooled.append(pooled)
+            text_pooled = torch.cat(text_pooled, dim=0)
+            
+            text_pooled = model.text_model.final_layer_norm(text_pooled)
+            text_pooled = model.text_projection(text_pooled)
+            text_pooled = F.normalize(text_pooled + eps, dim=-1)
+
+            seg_text_embeds = F.normalize(seg_text_embeds + eps, dim=-1)
+            
+            # Compute text alignment loss with cosine similarity directly
+            sim_text = text_pooled @ seg_text_embeds.t()  # removed logit_scale
+            text_diag = torch.diag(sim_text)
+            loss_text = mse_loss(text_diag, torch.ones_like(text_diag))
+            
+            # Total loss
+            loss = loss_org + 0.5 * loss_seg + loss_patch + loss_text
+            
+            epoch_loss += loss.item()
+            epoch_loss_org += loss_org.item()
+            epoch_loss_seg += loss_seg.item()
+            epoch_loss_patch += loss_patch.item()
+            epoch_loss_text += loss_text.item()
+            
+            fabric.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+            
+            if fabric.global_rank == 0:
+                wandb.log({
+                    "iter": iter,
+                    "lr": lr,
+                    "loss": loss.item(),
+                    "loss_org": loss_org.item(),
+                    "loss_seg": loss_seg.item(),
+                    "loss_patch": loss_patch.item(),
+                    "loss_text": loss_text.item(),
+                    "epoch": epoch,
+                    "progress": (iter / total_iter) * 100,
+                    "batch_size": args.batch_size,
+                    "logit_scale": model.logit_scale.exp().item(),
+                    "patch_similarity": patch_diag.mean().item(),  # average patch similarity
+                    "text_similarity": text_diag.mean().item(),    # average text similarity
+                })
+            
+            fabric.print(f"epoch {epoch} iter {iter} ({(iter/total_iter)*100:.4f}%) lr {lr:.6f} "
+                        f"loss {loss.item():.4f} (org: {loss_org.item():.4f}, seg: {loss_seg.item():.4f}, "
+                        f"patch: {loss_patch.item():.4f}, text: {loss_text.item():.4f} "
+                        f"patch_sim: {patch_diag.mean().item():.4f}, text_sim: {text_diag.mean().item():.4f})")
+            iter += 1
+            
+        # Calculate and log epoch averages
+        avg_epoch_loss = epoch_loss / len(train_loader)
+        avg_epoch_loss_org = epoch_loss_org / len(train_loader)
+        avg_epoch_loss_seg = epoch_loss_seg / len(train_loader)
+        avg_epoch_loss_patch = epoch_loss_patch / len(train_loader)
+        avg_epoch_loss_text = epoch_loss_text / len(train_loader)
+        
+        if fabric.global_rank == 0:
+            wandb.log({
+                "epoch": epoch,
+                "avg_epoch_loss": avg_epoch_loss,
+                "avg_epoch_loss_org": avg_epoch_loss_org,
+                "avg_epoch_loss_seg": avg_epoch_loss_seg,
+                "avg_epoch_loss_patch": avg_epoch_loss_patch,
+                "avg_epoch_loss_text": avg_epoch_loss_text,
+            })
+            
+        # Save model weights
+        save_path = os.path.join(args.output_dir, 
+                               f"GOAL_12_{os.path.splitext(os.path.basename(args.model))[0]}_"
+                               f"{os.path.splitext(os.path.basename(args.dataset))[0]}_{epoch+1}_{args.image_size}.pth")
+        
+        fabric.barrier()
+        if fabric.global_rank == 0:
+            model_state_dict = model.state_dict()
+            cpu_state_dict = {k: v.cpu() for k, v in model_state_dict.items()}
+            torch.save(cpu_state_dict, save_path)
+            fabric.print(f"Model saved to {save_path}")
+        fabric.barrier()
+        
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('CLIP Training', add_help=False)
+    parser.add_argument('--batch_size', default=16, type=int,
+                        help='Batch size per GPU')
+    parser.add_argument('--epochs', default=10, type=int)
+    parser.add_argument('--image_size', default=224, type=int)
+    parser.add_argument('--new_max_token', default=248, type=int)
+    parser.add_argument('--dataset', default='datasets/docci_segment_sim_bbox_del_org.json', type=str)
+    parser.add_argument('--model', default='openai/clip-vit-base-patch16', type=str)
+    parser.add_argument('--weight_decay', type=float, default=0.05) 
+    parser.add_argument('--init_lr', type=float, default=5e-6, metavar='LR') # originally 5e-6
+    parser.add_argument('--min_lr', type=float, default=0, metavar='LR')
+    parser.add_argument('--output_dir', default='finetune_out_SA_1B_100k_plus_docci/goal_bbox_local_token_align_batch16_only_max_pair_base16_patch16_real',
+                        help='path where to save, empty for no saving')
+    parser.add_argument('--save_interval', default=1, type=int)
+    parser.add_argument('--seed', default=0, type=int)
+    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--pin_mem', action='store_true',
+                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
+    parser.add_argument('--wandb_project', type=str, default='CLIP_Training', help='wandb project name')
+    parser.add_argument('--ckpt', type=str, default=None, help='path to checkpoint file')
+    parser.set_defaults(pin_mem=True)
+    parser.set_defaults(pin_mem=True)
+
+    # distributed training parameters
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='number of distributed processes')
+
+    return parser
+
+if __name__ == "__main__":
+    args = get_args_parser()
+    args = args.parse_args()
+    if args.output_dir:
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    main(args)

GOAL_github/mAP_goal_jointtest.py

@@ -0,0 +1,256 @@
+import json
+import argparse
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "1" 
+import sys
+from pathlib import Path
+import lightning as L
+import numpy as np
+import torch
+import transformers
+from torch.utils.data import DataLoader, Dataset, Subset
+from PIL import Image
+import torch.nn.functional as F
+from utils.func import *
+from utils.transforms import *
+
+# Hyperparameters
+micro_batch_size = 32
+devices = 1
+num_workers = 1
+
+class QueryLoader(Dataset):
+    def __init__(self, data_list, processor):
+        self.data_list = data_list
+        self.processor = processor
+        self.image_to_segs, self.org_images, self.filename_to_caption = self._create_mappings()
+
+    def _create_mappings(self):
+        image_to_segs = {}
+        org_images = set()
+        filename_to_caption = {}
+        for item in self.data_list:
+            filename = item['filename']
+            filename_to_caption[filename] = item['caption']
+            if 'segment_with_background' in filename:
+                org_filename = get_org_filename(filename)
+                if org_filename not in image_to_segs:
+                    image_to_segs[org_filename] = []
+                image_to_segs[org_filename].append(item)
+            else:
+                org_images.add(filename)
+                image_to_segs[filename] = []
+        return image_to_segs, org_images, filename_to_caption
+
+    def __len__(self):
+        return len(self.data_list)
+    
+    def _load_image(self, id: int):
+        return Image.open(self.data_list[id]["filename"]).convert("RGB")
+
+    def _load_target(self, id: int):
+        return self.data_list[id]["caption"]
+    
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+        image = self._load_image(idx)
+        caption = self._load_target(idx)
+        data = self.processor(images=image, text=caption, return_tensors="pt", truncation=True, padding="max_length", max_length=args.new_max_token)
+        return data.pixel_values[0], data.input_ids[0], self.data_list[idx]["filename"]
+
+def process_chunk(fabric, model, data_loader):
+    images = []
+    texts = []
+    filenames = []
+    for samples in data_loader:
+        image, text, filename = samples
+        with torch.no_grad():
+            x = model(pixel_values=image.to(fabric.device), input_ids=text.to(fabric.device))
+        x_i = F.normalize(x.image_embeds)
+        x_t = F.normalize(x.text_embeds)
+        images.append(x_i)
+        texts.append(x_t)
+        filenames.extend(filename)
+    return torch.cat(images), torch.cat(texts), filenames
+
+def compute_similarity(images, texts):
+    return torch.mm(images, texts.t())
+
+def compute_ap(ranks, relevant_items, k=None):
+    if not relevant_items:
+        return 0.0
+    score = 0.0
+    num_hits = 0.0
+    for i, item in enumerate(ranks[:k] if k else ranks):
+        if item in relevant_items:
+            num_hits += 1.0
+            score += num_hits / (i + 1.0)
+    return score / len(relevant_items)
+
+def get_org_filename(filename):
+    if 'segment_with_background' in filename:
+        parts = filename.split('/')
+        org_filename = parts[-2].split('_results')[0]
+        if not org_filename.endswith('.jpg'):
+            org_filename += '.jpg'
+        return org_filename
+    return filename
+
+def get_relevant_items(filename, all_filenames, image_to_segs, org_images, filename_to_caption):
+    org_filename = get_org_filename(filename)
+    
+    if filename in org_images:
+        # Query is an org image
+        relevant_items = [all_filenames.index(seg['filename']) for seg in image_to_segs[filename]]
+    else:
+        # Query is a seg image
+        relevant_items = []
+        if org_filename in all_filenames:
+            relevant_items.append(all_filenames.index(org_filename))
+        if filename in all_filenames:
+            relevant_items.append(all_filenames.index(filename))
+
+    return relevant_items
+
+def get_relevant_captions(filename, image_to_segs, org_images, filename_to_caption):
+    org_filename = get_org_filename(filename)
+    
+    if filename in org_images:
+        # Query is an org image
+        relevant_captions = [filename_to_caption[filename]]  # 원본 이미지의 캡션 추가
+        relevant_captions.extend([seg['caption'] for seg in image_to_segs[filename]])
+    else:
+        # Query is a seg image
+        relevant_captions = [filename_to_caption[filename]]
+        if org_filename in filename_to_caption:
+            relevant_captions.append(filename_to_caption[org_filename])
+
+    return relevant_captions
+
+def get_relevant_items_for_text(query_caption, all_filenames, image_to_segs, org_images, filename_to_caption):
+    relevant_items = []
+    for filename, caption in filename_to_caption.items():
+        if caption == query_caption:
+            if filename in org_images:
+                # 쿼리가 org 이미지의 캡션인 경우
+                relevant_items.append(all_filenames.index(filename))  # org 이미지
+                relevant_items.extend([all_filenames.index(seg['filename']) for seg in image_to_segs[filename]])  # seg 이미지들
+            else:
+                # 쿼리가 seg 이미지의 캡션인 경우
+                relevant_items.append(all_filenames.index(filename))  # seg 이미지
+                org_filename = get_org_filename(filename)
+                if org_filename in all_filenames:
+                    relevant_items.append(all_filenames.index(org_filename))  # org 이미지
+    return list(set(relevant_items))  # 중복 제거
+
+@torch.no_grad()
+def test(fabric: L.Fabric, model: torch.nn.Module, query_loader, k=None) -> torch.Tensor:
+    fabric.print("Testing ...")
+    
+    chunk_size = 5000
+    dataset_size = len(query_loader.dataset)
+    all_images = []
+    all_texts = []
+    all_filenames = []
+
+    for start_idx in range(0, dataset_size, chunk_size):
+        end_idx = min(start_idx + chunk_size, dataset_size)
+        chunk_dataset = Subset(query_loader.dataset, range(start_idx, end_idx))
+        chunk_loader = DataLoader(chunk_dataset, batch_size=query_loader.batch_size, shuffle=False, num_workers=query_loader.num_workers)
+        
+        chunk_images, chunk_texts, chunk_filenames = process_chunk(fabric, model, chunk_loader)
+        
+        all_images.append(chunk_images)
+        all_texts.append(chunk_texts)
+        all_filenames.extend(chunk_filenames)
+        
+        torch.cuda.empty_cache()
+
+    all_images = torch.cat(all_images)
+    all_texts = torch.cat(all_texts)
+
+    similarity = compute_similarity(all_images, all_texts)
+
+    image_to_segs = query_loader.dataset.image_to_segs
+    org_images = query_loader.dataset.org_images
+    filename_to_caption = query_loader.dataset.filename_to_caption
+    mAP_i2t = 0.0
+    mAP_t2i = 0.0
+
+    # Image to Text
+    for i, filename in enumerate(all_filenames):
+        relevant_captions = get_relevant_captions(filename, image_to_segs, org_images, filename_to_caption)
+
+        i2t_ranks = torch.argsort(similarity[i], descending=True).tolist()
+        i2t_relevant = [idx for idx, fn in enumerate(all_filenames) if filename_to_caption[fn] in relevant_captions]
+        mAP_i2t += compute_ap(i2t_ranks, i2t_relevant, k)
+
+    # Text to Image
+    unique_captions = set(filename_to_caption.values())
+    for caption in unique_captions:
+        relevant_items = get_relevant_items_for_text(caption, all_filenames, image_to_segs, org_images, filename_to_caption)
+        
+        caption_index = all_filenames.index(next(filename for filename, cap in filename_to_caption.items() if cap == caption))
+        t2i_ranks = torch.argsort(similarity[:, caption_index], descending=True).tolist()
+        mAP_t2i += compute_ap(t2i_ranks, relevant_items, k)
+
+    mAP_i2t /= len(all_filenames)
+    mAP_t2i /= len(unique_captions)
+
+    fabric.print(f"mAP@{k if k else 'all'} - Text-to-Image: {mAP_t2i:.4f} & Image-to-Text: {mAP_i2t:.4f}")
+
+
+def main(args):
+    fabric = L.Fabric(
+        accelerator="cuda", 
+        devices=devices,
+        precision="bf16-mixed"
+    )
+    fabric.launch()
+    fabric.seed_everything(1337 + fabric.global_rank)
+
+    if args.model == 'L-336':
+        args.model = 'openai/clip-vit-large-patch14-336'
+    elif args.model == 'L':
+        args.model = 'openai/clip-vit-large-patch14'
+    elif args.model == 'B':
+        args.model = 'openai/clip-vit-base-patch16'
+    elif args.model == 'G':
+        args.model = 'laion/CLIP-ViT-bigG-14-laion2B-39B-b160k'
+        
+    with fabric.device:
+        processor = transformers.AutoProcessor.from_pretrained(args.model)
+        model = transformers.AutoModel.from_pretrained(args.model).bfloat16()
+        longclip_pos_embeddings(model, args.new_max_token)
+        model.load_state_dict(torch.load(args.ckpt), strict=False)
+
+    if args.dataset == 'docci':
+        query_list = 'datasets/docci_test_joint_sim_max_1:1.json'
+    elif args.dataset == 'DCI':
+        query_list = 'datasets/DCI_test_joint_sim_max_1:1.json'
+
+    with open(query_list) as f:
+        query_list = json.loads(f.read())
+
+    args.query_list = query_list
+
+    query_dataset = QueryLoader(query_list, processor)
+    query_loader = DataLoader(query_dataset, num_workers=num_workers, batch_size=micro_batch_size, shuffle=False, drop_last=False, pin_memory=False)
+    query_loader = fabric.setup_dataloaders(query_loader)
+
+    model.eval().to(fabric.device)
+    test(fabric, model, query_loader, args.k)
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision("high")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, default='docci')
+    parser.add_argument('--new_max_token', default=248, type=int)
+    parser.add_argument("--model", type=str, default='B')
+    parser.add_argument("--ckpt", type=str, default='')
+    parser.add_argument("--k", type=int, default=10, help="Limit rank calculation to top K results. Use None for all ranks.")
+    args = parser.parse_args()
+
+    main(args)

GOAL_github/retrieval_goal.py

@@ -0,0 +1,171 @@
+import json
+import argparse
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+import sys
+from pathlib import Path
+import lightning as L
+import numpy as np
+import torch
+import transformers
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from PIL import Image
+import torch.nn.functional as F
+from utils.func import *
+from utils.transforms import *
+
+# Hyperparameters
+micro_batch_size = 32
+devices = 1
+num_workers = 1
+
+class QueryLoader(Dataset):
+    def __init__(self, data_list, processor):
+        self.data_list = data_list
+        self.processor = processor
+
+    def __len__(self):
+        return len(self.data_list)
+    
+    def _load_image(self, id: int):
+        return Image.open(self.data_list[id]["filename"]).convert("RGB")
+
+    def _load_target(self, id: int):
+        return self.data_list[id]["caption"]
+    
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+        image = self._load_image(idx)
+        caption = self._load_target(idx)
+        data = self.processor(images=image, text=caption, return_tensors="pt", truncation = True, padding = "max_length", max_length=args.new_max_token)
+    
+        return data.pixel_values[0], data.input_ids[0]
+    
+
+def main(args):
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # GPU 1번 인덱스만 사용
+    fabric = L.Fabric(
+        accelerator="cuda", 
+        devices=devices,
+        precision="bf16-mixed"  # "32"에서 "bf16-mixed"로 변경
+    )
+    fabric.launch()
+    fabric.seed_everything(1337 + fabric.global_rank)
+
+    if args.model=='L-336':
+        args.model = 'openai/clip-vit-large-patch14-336'
+    elif args.model=='L':
+        args.model = 'openai/clip-vit-large-patch14'
+    elif args.model=='B':
+        args.model = 'openai/clip-vit-base-patch16'
+    elif args.model=='G':
+        args.model = 'laion/CLIP-ViT-bigG-14-laion2B-39B-b160k'
+        
+    with fabric.device:
+        processor = transformers.AutoProcessor.from_pretrained(args.model)
+        model = transformers.AutoModel.from_pretrained(args.model).bfloat16()
+        longclip_pos_embeddings(model, args.new_max_token)
+        if args.ckpt:  # ckpt가 제공된 경우에만 로드
+            model.load_state_dict(torch.load(args.ckpt), strict=False)
+
+    if args.dataset == 'docci':
+        query_list = 'datasets/docci_test.json'
+    elif args.dataset =='coco':
+        query_list = 'datasets/coco_test.json'
+    elif args.dataset =='flickr30k':
+        query_list = 'datasets/flickr30k_test.json'    
+    elif args.dataset =='DCI':
+        query_list = 'datasets/DCI_test.json'
+    elif args.dataset =='urban':
+        query_list = 'datasets/urban_dataset_test.json'
+    elif args.dataset =='sharegpt4v':
+        query_list = 'datasets/sharegpt4v_test.json'
+
+    with open(query_list) as f:
+        query_list = json.loads(f.read())
+
+    args.query_list = query_list
+
+    query_dataset = QueryLoader(query_list, processor)
+    query_loader = DataLoader(query_dataset, num_workers=num_workers, batch_size=micro_batch_size, shuffle=False, drop_last=False, pin_memory=False)
+    query_loader = fabric.setup_dataloaders(query_loader)
+
+    model.eval().to(fabric.device)
+    test(fabric, model, query_loader)
+
+def compute_AP_and_recall_at_Ks(similarity, label_matrix, Ks):
+    # Sort gallery indices based on similarity
+    sorted_indices = torch.argsort(similarity, descending=True)
+    # Initialize results
+    results = {K: {'AP': 0.0, 'recall': 0.0, 'relevant_items': 0} for K in Ks}
+    total_relevant_items = label_matrix.sum().item()
+    for i, idx in enumerate(sorted_indices):
+        if label_matrix[idx]:
+            for K in Ks:
+                if i < K:
+                    results[K]['relevant_items'] += 1
+                    precision = results[K]['relevant_items'] / (i + 1)
+                    results[K]['AP'] += precision
+                    results[K]['recall'] = results[K]['relevant_items'] / total_relevant_items
+    for K in Ks:
+        results[K]['AP'] /= total_relevant_items
+    return results
+
+
+@torch.no_grad()
+def test(fabric: L.Fabric, model: torch.nn.Module, query_loader) -> torch.Tensor:
+    fabric.print("Testing ...")
+    
+    images = torch.tensor([], dtype=torch.float32).to(fabric.device)
+    texts = torch.tensor([], dtype=torch.float32).to(fabric.device)
+
+    for samples in query_loader:
+        image, text = samples
+
+        x = model(pixel_values=image, input_ids=text)
+
+        x_i = F.normalize(x.image_embeds)
+        x_t = F.normalize(x.text_embeds)
+
+        images = torch.cat((images,x_i), dim=0)
+        texts = torch.cat((texts,x_t), dim=0)
+
+    # Calculate cosine similarity
+    similarity = torch.mm(images, texts.t())
+    # Image to Text (I2T)
+    sorted_indices_i2t = torch.argsort(similarity, descending=True)
+    correct_indices = torch.arange(images.shape[0]).to(fabric.device)
+    ranks_i2t = (sorted_indices_i2t == correct_indices[:, None]).nonzero(as_tuple=True)[1]
+    # Text to Image (T2I)
+    sorted_indices_t2i = torch.argsort(similarity.t(), descending=True)
+    ranks_t2i = (sorted_indices_t2i == correct_indices[:, None]).nonzero(as_tuple=True)[1]
+    # Calculate recall at different ranks for I2T
+    recall_i2t_1 = (ranks_i2t < 1).float().mean().item() * 100
+    recall_i2t_5 = (ranks_i2t < 5).float().mean().item() * 100
+    recall_i2t_25 = (ranks_i2t < 25).float().mean().item() * 100
+    recall_i2t_50 = (ranks_i2t < 50).float().mean().item() * 100
+    # Calculate recall at different ranks for T2I
+    recall_t2i_1 = (ranks_t2i < 1).float().mean().item() * 100
+    recall_t2i_5 = (ranks_t2i < 5).float().mean().item() * 100
+    recall_t2i_25 = (ranks_t2i < 25).float().mean().item() * 100
+    recall_t2i_50 = (ranks_t2i < 50).float().mean().item() * 100
+    # Print recall percentages for T2I
+    fabric.print(f"Text-to-Image: {recall_t2i_1:.2f} & {recall_t2i_5:.2f} & {recall_t2i_25:.2f} & {recall_t2i_50:.2f}")
+    # Print recall percentages for I2T
+    fabric.print(f"Image-to-Text: {recall_i2t_1:.2f} & {recall_i2t_5:.2f} & {recall_i2t_25:.2f} & {recall_i2t_50:.2f}")
+
+
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision("high")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, default='urban')
+    parser.add_argument('--new_max_token', default=248, type=int)
+    parser.add_argument("--model", type=str, default='B')
+    parser.add_argument("--ckpt", type=str, default='')
+    args = parser.parse_args()
+
+    main(args)

GOAL_github/utils/func.py

@@ -0,0 +1,106 @@
+import torch
+import cv2
+from PIL import Image
+import numpy as np
+
+def interpolate_pos_embeddings(model, new_image_size):
+    vision_model = model.vision_model
+    patch_size = vision_model.config.patch_size
+    num_patches = (new_image_size // patch_size) ** 2 + 1
+    # Extract and interpolate positional embeddings
+    pos_embeddings = vision_model.embeddings.position_embedding.weight
+    pos_embeddings = pos_embeddings.unsqueeze(0).permute(0, 2, 1)  # Convert to 1xCxN format
+    pos_embeddings = torch.nn.functional.interpolate(
+        pos_embeddings, size=(num_patches), mode='nearest'
+    ).squeeze(0).permute(1, 0)  # Convert back to NxC format
+    pos_embeddings = pos_embeddings.contiguous()  # Ensure contiguous
+    vision_model.embeddings.position_embedding.weight = torch.nn.Parameter(pos_embeddings)
+    # Set position_ids
+    if hasattr(vision_model.embeddings, 'position_ids'):
+        vision_model.embeddings.position_ids = torch.arange(0, num_patches).unsqueeze(0)
+    else:
+        vision_model.register_buffer('position_ids', torch.arange(0, num_patches).unsqueeze(0))
+
+def interpolate_text_pos_embeddings(model, new_max_token):
+    text_model = model.text_model
+    # Extract and interpolate positional embeddings
+    pos_embeddings = text_model.embeddings.position_embedding.weight
+    pos_embeddings = pos_embeddings.unsqueeze(0).permute(0, 2, 1)  # Convert to 1xCxN format
+    
+    # Interpolate the position embeddings to the new maximum token length
+    pos_embeddings = torch.nn.functional.interpolate(
+        pos_embeddings, size=(new_max_token), mode='nearest'
+    ).squeeze(0).permute(1, 0)  # Convert back to NxC format
+    
+    pos_embeddings = pos_embeddings.contiguous()  # Ensure contiguous
+    text_model.embeddings.position_embedding.weight = torch.nn.Parameter(pos_embeddings)
+    
+    # Set position_ids if the model uses them
+    if hasattr(text_model.embeddings, 'position_ids'):
+        text_model.embeddings.position_ids = torch.arange(0, new_max_token).unsqueeze(0)
+    else:
+        text_model.register_buffer('position_ids', torch.arange(0, new_max_token).unsqueeze(0))
+
+def longclip_pos_embeddings(model, new_max_token):
+    text_model = model.text_model
+    # Extract positional embeddings
+    pos_embeddings_pre = text_model.embeddings.position_embedding.weight
+    length, dim = pos_embeddings_pre.shape
+    keep_len = 20
+    new_length = 4*length - 3*keep_len
+    if new_length < new_max_token:
+        raise ValueError("new_max_token is too large")
+    pos_embeddings_new = torch.zeros([new_max_token, dim], dtype=pos_embeddings_pre.dtype)
+    for i in range(keep_len):
+        pos_embeddings_new[i] = pos_embeddings_pre[i]
+    for i in range(length-1-keep_len):
+        pos_embeddings_new[4*i + keep_len] = pos_embeddings_pre[i + keep_len]
+        pos_embeddings_new[4*i + 1 + keep_len] = 3*pos_embeddings_pre[i + keep_len]/4 + 1*pos_embeddings_pre[i+1+keep_len]/4
+        pos_embeddings_new[4*i + 2+keep_len] = 2*pos_embeddings_pre[i+keep_len]/4 + 2*pos_embeddings_pre[i+1+keep_len]/4
+        pos_embeddings_new[4*i + 3+keep_len] = 1*pos_embeddings_pre[i+keep_len]/4 + 3*pos_embeddings_pre[i+1+keep_len]/4
+    pos_embeddings_new[4*length -3*keep_len - 4] = pos_embeddings_pre[length-1] + 0*(pos_embeddings_pre[length-1] - pos_embeddings_pre[length-2])/4
+    pos_embeddings_new[4*length -3*keep_len - 3] = pos_embeddings_pre[length-1] + 1*(pos_embeddings_pre[length-1] - pos_embeddings_pre[length-2])/4
+    pos_embeddings_new[4*length -3*keep_len - 2] = pos_embeddings_pre[length-1] + 2*(pos_embeddings_pre[length-1] - pos_embeddings_pre[length-2])/4
+    pos_embeddings_new[4*length -3*keep_len - 1] = pos_embeddings_pre[length-1] + 3*(pos_embeddings_pre[length-1] - pos_embeddings_pre[length-2])/4
+    text_model.embeddings.position_embedding.weight = torch.nn.Parameter(pos_embeddings_new)
+    # Set position_ids if the model uses them
+    if hasattr(text_model.embeddings, 'position_ids'):
+        text_model.embeddings.position_ids = torch.arange(0, new_max_token).unsqueeze(0)
+    else:
+        text_model.register_buffer('position_ids', torch.arange(0, new_max_token).unsqueeze(0))
+
+
+def average_pool(last_hidden_states, attention_mask):
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+def last_token_pool(last_hidden_states, attention_mask):
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+def batch_align(fabric, x):
+    x = fabric.all_gather(x, sync_grads=True)
+    return x.view(x.shape[0]*x.shape[1], -1)
+
+cls_criterion = torch.nn.CrossEntropyLoss()
+
+def clip_loss(logits):
+    gt = torch.arange(len(logits),dtype=torch.long, device=logits.device)
+    return (cls_criterion(logits, gt) + cls_criterion(logits.t(), gt))/2.0
+
+def print_trainable_parameters(fabric, model):
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    fabric.print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param:.2f}"
+    )
+    fabric.print('Memory load of model: {} bytes'.format(torch.cuda.memory_allocated()))

GOAL_github/utils/randaugment.py

@@ -0,0 +1,349 @@
+# Code from salesforce/BLIP
+
+import cv2
+import numpy as np
+
+
+## aug functions
+def identity_func(img):
+    return img
+
+
+def autocontrast_func(img, cutoff=0):
+    """
+    same output as PIL.ImageOps.autocontrast
+    """
+    n_bins = 256
+
+    def tune_channel(ch):
+        n = ch.size
+        cut = cutoff * n // 100
+        if cut == 0:
+            high, low = ch.max(), ch.min()
+        else:
+            hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+            low = np.argwhere(np.cumsum(hist) > cut)
+            low = 0 if low.shape[0] == 0 else low[0]
+            high = np.argwhere(np.cumsum(hist[::-1]) > cut)
+            high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0]
+        if high <= low:
+            table = np.arange(n_bins)
+        else:
+            scale = (n_bins - 1) / (high - low)
+            offset = np.multiply(low, -scale)
+            table = np.arange(n_bins) * scale + offset
+            table[table < 0] = 0
+            table[table > n_bins - 1] = n_bins - 1
+        table = table.clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def equalize_func(img):
+    """
+    same output as PIL.ImageOps.equalize
+    PIL's implementation is different from cv2.equalize
+    """
+    n_bins = 256
+
+    def tune_channel(ch):
+        hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+        non_zero_hist = hist[hist != 0].reshape(-1)
+        step = np.sum(non_zero_hist[:-1]) // (n_bins - 1)
+        if step == 0:
+            return ch
+        n = np.empty_like(hist)
+        n[0] = step // 2
+        n[1:] = hist[:-1]
+        table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def rotate_func(img, degree, fill=(0, 0, 0)):
+    """
+    like PIL, rotate by degree, not radians
+    """
+    H, W = img.shape[0], img.shape[1]
+    center = W / 2, H / 2
+    M = cv2.getRotationMatrix2D(center, degree, 1)
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill)
+    return out
+
+
+def solarize_func(img, thresh=128):
+    """
+    same output as PIL.ImageOps.posterize
+    """
+    table = np.array([el if el < thresh else 255 - el for el in range(256)])
+    table = table.clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def color_func(img, factor):
+    """
+    same output as PIL.ImageEnhance.Color
+    """
+    ## implementation according to PIL definition, quite slow
+    #  degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
+    #  out = blend(degenerate, img, factor)
+    #  M = (
+    #      np.eye(3) * factor
+    #      + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor)
+    #  )[np.newaxis, np.newaxis, :]
+    M = np.array(
+        [[0.886, -0.114, -0.114], [-0.587, 0.413, -0.587], [-0.299, -0.299, 0.701]],
+        dtype=np.float32,
+    ) * factor + np.array([[0.114], [0.587], [0.299]], dtype=np.float32)
+    out = np.matmul(img, M).clip(0, 255).astype(np.uint8)
+    return out
+
+
+def contrast_func(img, factor):
+    """
+    same output as PIL.ImageEnhance.Contrast
+    """
+    mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299]))
+    table = (
+        np.array([(el - mean) * factor + mean for el in range(256)])
+        .clip(0, 255)
+        .astype(np.uint8)
+    )
+    out = table[img]
+    return out
+
+
+def brightness_func(img, factor):
+    """
+    same output as PIL.ImageEnhance.Contrast
+    """
+    table = (np.arange(256, dtype=np.float32) * factor).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def sharpness_func(img, factor):
+    """
+    The differences the this result and PIL are all on the 4 boundaries, the center
+    areas are same
+    """
+    kernel = np.ones((3, 3), dtype=np.float32)
+    kernel[1][1] = 5
+    kernel /= 13
+    degenerate = cv2.filter2D(img, -1, kernel)
+    if factor == 0.0:
+        out = degenerate
+    elif factor == 1.0:
+        out = img
+    else:
+        out = img.astype(np.float32)
+        degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :]
+        out[1:-1, 1:-1, :] = degenerate + factor * (out[1:-1, 1:-1, :] - degenerate)
+        out = out.astype(np.uint8)
+    return out
+
+
+def shear_x_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.array([[1, factor, 0], [0, 1, 0]], dtype=np.float32)
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def translate_x_func(img, offset, fill=(0, 0, 0)):
+    """
+    same output as PIL.Image.transform
+    """
+    H, W = img.shape[0], img.shape[1]
+    M = np.array([[1, 0, -offset], [0, 1, 0]], dtype=np.float32)
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def translate_y_func(img, offset, fill=(0, 0, 0)):
+    """
+    same output as PIL.Image.transform
+    """
+    H, W = img.shape[0], img.shape[1]
+    M = np.array([[1, 0, 0], [0, 1, -offset]], dtype=np.float32)
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def posterize_func(img, bits):
+    """
+    same output as PIL.ImageOps.posterize
+    """
+    out = np.bitwise_and(img, np.uint8(255 << (8 - bits)))
+    return out
+
+
+def shear_y_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.array([[1, 0, 0], [factor, 1, 0]], dtype=np.float32)
+    out = cv2.warpAffine(
+        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR
+    ).astype(np.uint8)
+    return out
+
+
+def cutout_func(img, pad_size, replace=(0, 0, 0)):
+    replace = np.array(replace, dtype=np.uint8)
+    H, W = img.shape[0], img.shape[1]
+    rh, rw = np.random.random(2)
+    pad_size = pad_size // 2
+    ch, cw = int(rh * H), int(rw * W)
+    x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)
+    y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)
+    out = img.copy()
+    out[x1:x2, y1:y2, :] = replace
+    return out
+
+
+### level to args
+def enhance_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        return ((level / MAX_LEVEL) * 1.8 + 0.1,)
+
+    return level_to_args
+
+
+def shear_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 0.3
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def translate_level_to_args(translate_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * float(translate_const)
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * cutout_const)
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def solarize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 256)
+        return (level,)
+
+    return level_to_args
+
+
+def none_level_to_args(level):
+    return ()
+
+
+def posterize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 4)
+        return (level,)
+
+    return level_to_args
+
+
+def rotate_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 30
+        if np.random.random() < 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+func_dict = {
+    "Identity": identity_func,
+    "AutoContrast": autocontrast_func,
+    "Equalize": equalize_func,
+    "Rotate": rotate_func,
+    "Solarize": solarize_func,
+    "Color": color_func,
+    "Contrast": contrast_func,
+    "Brightness": brightness_func,
+    "Sharpness": sharpness_func,
+    "ShearX": shear_x_func,
+    "TranslateX": translate_x_func,
+    "TranslateY": translate_y_func,
+    "Posterize": posterize_func,
+    "ShearY": shear_y_func,
+}
+
+translate_const = 10
+MAX_LEVEL = 10
+replace_value = (128, 128, 128)
+arg_dict = {
+    "Identity": none_level_to_args,
+    "AutoContrast": none_level_to_args,
+    "Equalize": none_level_to_args,
+    "Rotate": rotate_level_to_args(MAX_LEVEL, replace_value),
+    "Solarize": solarize_level_to_args(MAX_LEVEL),
+    "Color": enhance_level_to_args(MAX_LEVEL),
+    "Contrast": enhance_level_to_args(MAX_LEVEL),
+    "Brightness": enhance_level_to_args(MAX_LEVEL),
+    "Sharpness": enhance_level_to_args(MAX_LEVEL),
+    "ShearX": shear_level_to_args(MAX_LEVEL, replace_value),
+    "TranslateX": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),
+    "TranslateY": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),
+    "Posterize": posterize_level_to_args(MAX_LEVEL),
+    "ShearY": shear_level_to_args(MAX_LEVEL, replace_value),
+}
+
+
+class RandomAugment(object):
+    def __init__(self, N=2, M=10, isPIL=False, augs=[]):
+        self.N = N
+        self.M = M
+        self.isPIL = isPIL
+        if augs:
+            self.augs = augs
+        else:
+            self.augs = list(arg_dict.keys())
+
+    def get_random_ops(self):
+        sampled_ops = np.random.choice(self.augs, self.N)
+        return [(op, 0.5, self.M) for op in sampled_ops]
+
+    def __call__(self, img):
+        if self.isPIL:
+            img = np.array(img)
+        ops = self.get_random_ops()
+        for name, prob, level in ops:
+            if np.random.random() > prob:
+                continue
+            args = arg_dict[name](level)
+            img = func_dict[name](img, *args)
+        return img
+
+
+if __name__ == "__main__":
+    a = RandomAugment()
+    img = np.random.randn(32, 32, 3)
+    a(img)

GOAL_github/utils/transforms.py

@@ -0,0 +1,130 @@
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+from torchvision.transforms import Compose, CenterCrop, ToTensor, Normalize, Resize
+from utils.randaugment import RandomAugment
+import torch
+import torchvision.transforms.functional as FT
+import math
+import torch.nn.functional as F
+
+normalize = transforms.Normalize(
+    (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
+)
+
+class transform_train:
+    def __init__(self, image_size=384, min_scale=0.5):
+        self.transform = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(
+                    image_size,
+                    scale=(min_scale, 1.0),
+                    interpolation=InterpolationMode.BICUBIC,
+                ),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+
+    def __call__(self, img):
+        return self.transform(img)
+
+
+# class transform_train:
+#     def __init__(self, image_size=384, min_scale=0.5):
+#         self.transform = transforms.Compose(
+#             [
+#                 transforms.RandomResizedCrop(
+#                     image_size,
+#                     scale=(min_scale, 1.0),
+#                     interpolation=InterpolationMode.BICUBIC,
+#                 ),
+#                 transforms.RandomHorizontalFlip(),
+#                 RandomAugment(
+#                     2,
+#                     5,
+#                     isPIL=True,
+#                     augs=[
+#                         "Identity",
+#                         "AutoContrast",
+#                         "Brightness",
+#                         "Sharpness",
+#                         "Equalize",
+#                         "ShearX",
+#                         "ShearY",
+#                         "TranslateX",
+#                         "TranslateY",
+#                         "Rotate",
+#                     ],
+#                 ),
+#                 transforms.ToTensor(),
+#                 normalize,
+#             ]
+#         )
+
+#     def __call__(self, img):
+#         return self.transform(img)
+
+
+class transform_test(transforms.Compose):
+    def __init__(self, image_size=384):
+        self.transform = transforms.Compose(
+            [
+                transforms.Resize(
+                    (image_size, image_size),
+                    interpolation=InterpolationMode.BICUBIC,
+                ),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+
+    def __call__(self, img):
+        return self.transform(img)
+
+
+
+class TargetPad:
+    """
+    If an image aspect ratio is above a target ratio, pad the image to match such target ratio.
+    For more details see Baldrati et al. 'Effective conditioned and composed image retrieval combining clip-based features.' Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (2022).
+    """
+
+    def __init__(self, target_ratio: float, size: int):
+        """
+        :param target_ratio: target ratio
+        :param size: preprocessing output dimension
+        """
+        self.size = size
+        self.target_ratio = target_ratio
+
+    def __call__(self, image):
+        w, h = image.size
+        actual_ratio = max(w, h) / min(w, h)
+        if actual_ratio < self.target_ratio:  # check if the ratio is above or below the target ratio
+            return image
+        scaled_max_wh = max(w, h) / self.target_ratio  # rescale the pad to match the target ratio
+        hp = max(int((scaled_max_wh - w) / 2), 0)
+        vp = max(int((scaled_max_wh - h) / 2), 0)
+        padding = [hp, vp, hp, vp]
+        return FT.pad(image, padding, 0, 'constant')
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+def targetpad_transform(target_ratio: float, dim: int) -> torch.Tensor:
+    """
+    CLIP-like preprocessing transform computed after using TargetPad pad
+    :param target_ratio: target ratio for TargetPad
+    :param dim: image output dimension
+    :return: CLIP-like torchvision Compose transform
+    """
+    return Compose([
+        TargetPad(target_ratio, dim),
+        Resize(dim, interpolation=InterpolationMode.BICUBIC),
+        CenterCrop(dim),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])

GOAL_github/visualization/visualization_attentionmap_longtestset.py

@@ -0,0 +1,188 @@
+import json
+import argparse
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+import sys
+from pathlib import Path
+import lightning as L
+import numpy as np
+import torch
+import transformers
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from PIL import Image
+import torch.nn.functional as F
+from utils.func import *
+from utils.transforms import *
+import matplotlib.pyplot as plt
+import cv2
+
+# Hyperparameters
+micro_batch_size = 1
+devices = 1
+num_workers = 1
+
+class ImageAttentionLoader(Dataset):
+    def __init__(self, image_path, processor):
+        self.image_path = image_path
+        self.processor = processor
+
+    def __len__(self):
+        return 1
+    
+    def _load_image(self):
+        return Image.open(self.image_path).convert("RGB")
+    
+    def __getitem__(self, idx):
+        image = self._load_image()
+        self.original_size = image.size
+        # Process image without return_tensors
+        data = self.processor(images=image)
+        # Convert to tensor manually
+        pixel_values = torch.tensor(data['pixel_values'])
+        # Ensure correct shape [C, H, W]
+        if len(pixel_values.shape) > 3:
+            pixel_values = pixel_values.squeeze()
+        return pixel_values, np.array(image)
+
+def get_attention_maps(model, image_input):
+    print("Input shape before processing:", image_input.shape)
+    
+    # Ensure correct shape and device
+    if len(image_input.shape) == 3:
+        image_input = image_input.unsqueeze(0)  # Add batch dimension
+    
+    image_input = image_input.to(model.device)
+    print("Input shape after processing:", image_input.shape)
+    
+    outputs = model.vision_model(pixel_values=image_input, output_attentions=True)
+    
+    # 마지막 레이어의 어텐션 맵 가져오기
+    attention_maps = outputs.attentions[-1]
+    print("Attention maps shape:", attention_maps.shape)
+    
+    # 어텐션 맵의 크기 조정 (헤드 평균)
+    attention_maps = attention_maps.mean(dim=1)
+    
+    # CLS 토큰을 제외한 패치들의 어텐션만 사용
+    attention_maps = attention_maps[:, 1:, 1:]
+    
+    # CPU로 이동하고 float32로 변환
+    attention_maps = attention_maps.squeeze().cpu().float()
+    print("Final attention maps shape:", attention_maps.shape)
+    
+    return attention_maps
+
+def visualize_attention(image, attention_maps, output_path, model_name):
+    # 이미지 패치 크기 계산
+    if 'patch14' in model_name:
+        patch_size = 14
+    elif 'patch32' in model_name:
+        patch_size = 32
+    else:
+        patch_size = 16
+    
+    # 어텐션 맵 크기 조정
+    h = int(np.sqrt(attention_maps.shape[0]))
+    attention_maps = attention_maps.reshape(h, h)
+    
+    # 어텐션 맵을 원본 이미지 크기로 업샘플링
+    attention_maps = F.interpolate(
+        attention_maps.unsqueeze(0).unsqueeze(0),
+        size=image.shape[:2],
+        mode='bicubic'
+    ).squeeze().numpy()
+    
+    # 정규화
+    attention_maps = (attention_maps - attention_maps.min()) / (attention_maps.max() - attention_maps.min())
+    
+    # 시각화
+    plt.figure(figsize=(15, 5))
+    
+    # 원본 이미지
+    plt.subplot(1, 3, 1)
+    plt.imshow(image)
+    plt.title('Original Image')
+    plt.axis('off')
+    
+    # 어텐션 맵
+    plt.subplot(1, 3, 2)
+    plt.imshow(attention_maps, cmap='jet')
+    plt.title('Self-Attention Map')
+    plt.axis('off')
+    
+    # 오버레이
+    plt.subplot(1, 3, 3)
+    plt.imshow(image)
+    plt.imshow(attention_maps, cmap='jet', alpha=0.5)
+    plt.title('Overlay')
+    plt.axis('off')
+    
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    plt.close()
+
+def main(args):
+    os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+    fabric = L.Fabric(
+        accelerator="cuda",
+        devices=devices,
+        precision="bf16-mixed"
+    )
+    fabric.launch()
+    fabric.seed_everything(1337 + fabric.global_rank)
+
+    if args.model=='L-336':
+        args.model = 'openai/clip-vit-large-patch14-336'
+    elif args.model=='L':
+        args.model = 'openai/clip-vit-large-patch14'
+    elif args.model=='B':
+        args.model = 'openai/clip-vit-base-patch32'
+    elif args.model=='G':
+        args.model = 'laion/CLIP-ViT-bigG-14-laion2B-39B-b160k'
+        
+    with fabric.device:
+        processor = transformers.AutoProcessor.from_pretrained(args.model)
+        model = transformers.AutoModel.from_pretrained(args.model, output_attentions=True).bfloat16()
+        longclip_pos_embeddings(model, args.new_max_token)
+        if args.ckpt:
+            model.load_state_dict(torch.load(args.ckpt), strict=False)
+
+    # 데이터 로더 생성
+    dataset = ImageAttentionLoader(args.image_path, processor)
+    dataloader = DataLoader(dataset, batch_size=micro_batch_size, shuffle=False)
+
+    # 모델 평가 모드로 설정
+    model.eval()
+    model = model.to(fabric.device)
+    
+    # 어텐션 맵 생성 및 시각화
+    with torch.no_grad():
+        for batch in dataloader:
+            image_input, original_image = batch
+            print("Original input shape:", image_input.shape)
+            attention_maps = get_attention_maps(model, image_input)
+            
+            # 모든 패치의 평균 어텐션 점수 계산
+            avg_attention = attention_maps.mean(dim=0)
+            
+            visualize_attention(
+                original_image[0],
+                avg_attention,
+                args.output_path,
+                args.model
+            )
+            break
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision("high")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_path", type=str, required=True, help="Path to input image")
+    parser.add_argument("--output_path", type=str, required=True, help="Path to save visualization")
+    parser.add_argument('--new_max_token', default=248, type=int)
+    parser.add_argument("--model", type=str, default='L')
+    parser.add_argument("--ckpt", type=str, default='')
+    args = parser.parse_args()
+
+    main(args)

GOAL_github/visualization/visualization_retreival.py

@@ -0,0 +1,313 @@
+import json
+import argparse
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+import sys
+from pathlib import Path
+import lightning as L
+import numpy as np
+import torch
+import transformers
+from torch.utils.data import DataLoader, Dataset
+from PIL import Image
+import torch.nn.functional as F
+from utils.func import *
+import matplotlib.pyplot as plt
+import cv2
+import textwrap
+
+def collate_fn(batch):
+    # 배치에서 최대 길이 찾기
+    max_text_length = max(len(item['input_ids']) for item in batch)
+    
+    # 배치의 모든 항목을 담을 리스트
+    pixel_values = []
+    input_ids = []
+    attention_masks = []
+    image_paths = []
+    captions = []
+    
+    for item in batch:
+        # 이미지 처리
+        pixel_values.append(item['pixel_values'])
+        
+        # 텍스트 패딩
+        curr_len = len(item['input_ids'])
+        padding_len = max_text_length - curr_len
+        
+        padded_ids = torch.cat([
+            item['input_ids'],
+            torch.zeros(padding_len, dtype=torch.long)
+        ])
+        padded_mask = torch.cat([
+            item['attention_mask'],
+            torch.zeros(padding_len, dtype=torch.long)
+        ])
+        
+        input_ids.append(padded_ids)
+        attention_masks.append(padded_mask)
+        image_paths.append(item['image_path'])
+        captions.append(item['caption'])
+    
+    # 배치로 만들기
+    return {
+        'pixel_values': torch.stack(pixel_values),
+        'input_ids': torch.stack(input_ids),
+        'attention_mask': torch.stack(attention_masks),
+        'image_path': image_paths,
+        'caption': captions
+    }
+
+class JsonGalleryDataset(Dataset):
+    def __init__(self, json_path, processor, max_length=248):
+        with open(json_path, 'r', encoding='utf-8') as f:
+            self.data = json.load(f)
+        self.processor = processor
+        self.max_length = max_length
+        
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        image_path = item['filename']
+        caption = item['caption']
+        
+        try:
+            image = Image.open(image_path).convert("RGB")
+            # 이미지와 텍스트 모두 처리
+            processed_image = self.processor(images=image, return_tensors="pt")
+            processed_text = self.processor(
+                text=caption,
+                return_tensors="pt",
+                padding='max_length',
+                max_length=self.max_length,
+                truncation=True
+            )
+            
+            return {
+                'pixel_values': processed_image['pixel_values'].squeeze(0),
+                'input_ids': processed_text['input_ids'].squeeze(0),
+                'attention_mask': processed_text['attention_mask'].squeeze(0),
+                'image_path': image_path,
+                'caption': caption
+            }
+        except Exception as e:
+            print(f"Error loading image {image_path}: {e}")
+            return self.__getitem__(0)
+
+def compute_similarities(model, query_features, gallery_features):
+    query_features = F.normalize(query_features, dim=-1)
+    gallery_features = F.normalize(gallery_features, dim=-1)
+    similarities = torch.mm(query_features, gallery_features.t())
+    return similarities
+
+def process_query(model, processor, query, device, is_image=False, max_length=248):
+    with torch.no_grad():
+        if is_image:
+            image = Image.open(query).convert("RGB")
+            inputs = processor(images=image, return_tensors="pt")
+            pixel_values = inputs['pixel_values'].to(device)
+            features = model.get_image_features(pixel_values)
+        else:
+            inputs = processor(
+                text=query,
+                return_tensors="pt",
+                padding='max_length',
+                max_length=max_length,
+                truncation=True
+            )
+            input_ids = inputs['input_ids'].to(device)
+            attention_mask = inputs['attention_mask'].to(device)
+            features = model.get_text_features(input_ids, attention_mask)
+    return features
+
+def visualize_results(query, results, output_path, is_image_query=False):
+    # 하나의 결과만 사용
+    n_results = 1
+    plt.rcParams['figure.facecolor'] = 'white'
+    plt.rcParams['axes.facecolor'] = 'white'
+    fig, axes = plt.subplots(1, 1 + n_results, figsize=(10, 4))  # figure 크기 조정
+    
+    if is_image_query:
+        # 이미지 쿼리인 경우
+        query_img = Image.open(query)
+        axes[0].imshow(query_img)
+        axes[0].set_title("Image query", fontsize=12, pad=10)
+        
+        # 결과 표시 (캡션만)
+        img_path, similarity, caption = results[0]  # 첫 번째 결과만 사용
+        axes[1].set_facecolor('white')
+        wrapped_text = textwrap.fill(caption, width=40)
+        
+        # 텍스트 박스 없이 텍스트만 표시
+        axes[1].text(0.5, 0.5, wrapped_text, 
+                     horizontalalignment='center',
+                     verticalalignment='center',
+                     transform=axes[1].transAxes,
+                     wrap=True,
+                     fontsize=10)
+        # axes[1].set_title("Recall@1", fontsize=12, pad=10)
+            
+    else:
+        # 텍스트 쿼리인 경우
+        axes[0].set_facecolor('white')
+        wrapped_text = textwrap.fill(query, width=40)
+        
+        # 텍스트 박스 없이 텍스트만 표시
+        axes[0].text(0.5, 0.5, wrapped_text, 
+                    horizontalalignment='center',
+                    verticalalignment='center',
+                    transform=axes[0].transAxes,
+                    wrap=True,
+                    fontsize=10)
+        axes[0].set_title("Text query", fontsize=12, pad=10)
+        
+        # 결과 표시 (이미지만)
+        img_path, similarity, caption = results[0]  # 첫 번째 결과만 사용
+        img = Image.open(img_path)
+        axes[1].imshow(img)
+        # axes[1].set_title("Recall@1", fontsize=12, pad=10)
+    
+    # 모든 subplot의 공통 스타일 설정
+    for ax in axes:
+        ax.axis('off')
+        ax.set_xticks([])
+        ax.set_yticks([])
+        ax.spines['top'].set_visible(True)
+        ax.spines['right'].set_visible(True)
+        ax.spines['bottom'].set_visible(True)
+        ax.spines['left'].set_visible(True)
+    
+    plt.tight_layout()
+    plt.subplots_adjust(wspace=0)
+    
+    plt.savefig(output_path, 
+                dpi=300, 
+                bbox_inches='tight',
+                facecolor='white',
+                edgecolor='none',
+                pad_inches=0.2)
+    plt.close()
+    
+def main(args):
+    fabric = L.Fabric(
+        accelerator="cuda", 
+        devices=1,
+        precision="bf16-mixed"
+    )
+    fabric.launch()
+    fabric.seed_everything(1337 + fabric.global_rank)
+
+    # 모델 설정
+    if args.model == 'L-336':
+        model_name = 'openai/clip-vit-large-patch14-336'
+    elif args.model == 'L':
+        model_name = 'openai/clip-vit-large-patch14'
+    elif args.model == 'B':
+        model_name = 'openai/clip-vit-base-patch32'
+    elif args.model == 'G':
+        model_name = 'laion/CLIP-ViT-bigG-14-laion2B-39B-b160k'
+    
+    # 모델과 프로세서 로드
+    processor = transformers.CLIPProcessor.from_pretrained(model_name)
+    model = transformers.CLIPModel.from_pretrained(model_name).bfloat16()
+    
+    # 먼저 position embedding 확장
+    longclip_pos_embeddings(model, args.new_max_token)
+    
+    # 그 다음 가중치 로드
+    if args.ckpt:
+        state_dict = torch.load(args.ckpt, weights_only=True)
+        model.load_state_dict(state_dict, strict=False)
+    
+    model = model.to(fabric.device)
+    model.eval()
+
+    # 갤러리 데이터셋 준비
+    dataset = JsonGalleryDataset(args.gallery_json, processor, args.new_max_token)
+    dataloader = DataLoader(
+        dataset, 
+        batch_size=32, 
+        shuffle=False, 
+        num_workers=4,
+        collate_fn=collate_fn
+    )
+
+    # 쿼리가 이미지인지 텍스트인지 확인
+    is_image_query = args.query.endswith(('.jpg', '.jpeg', '.png'))
+    
+    # 갤러리 특징 추출
+    gallery_features = []
+    gallery_items = []
+    
+    with torch.no_grad():
+        for batch in dataloader:
+            if is_image_query:
+                # 이미지 쿼리의 경우 텍스트 특징 추출
+                input_ids = batch['input_ids'].to(fabric.device)
+                attention_mask = batch['attention_mask'].to(fabric.device)
+                features = model.get_text_features(input_ids, attention_mask)
+            else:
+                # 텍스트 쿼리의 경우 이미지 특징 추출
+                pixel_values = batch['pixel_values'].to(fabric.device)
+                features = model.get_image_features(pixel_values)
+            
+            gallery_features.append(features)
+            
+            for path, caption in zip(batch['image_path'], batch['caption']):
+                gallery_items.append({
+                    'image_path': path,
+                    'caption': caption
+                })
+    
+    gallery_features = torch.cat(gallery_features, dim=0)
+
+    # 쿼리 특징 추출
+    query_features = process_query(model, processor, args.query, fabric.device, is_image_query, args.new_max_token)
+
+    # 유사도 계산 및 상위 1개 결과만 가져오기
+    similarities = compute_similarities(model, query_features, gallery_features)
+    top_k = 1  # 하나의 결과만 가져오도록 변경
+    top_k_similarities, top_k_indices = torch.topk(similarities[0].float(), k=top_k)
+    
+    # 결과 정리
+    results = []
+    top_k_similarities_np = top_k_similarities.cpu().numpy()
+    top_k_indices_np = top_k_indices.cpu().numpy()
+    
+    for sim, idx in zip(top_k_similarities_np, top_k_indices_np):
+        item = gallery_items[idx]
+        results.append((
+            item['image_path'],
+            sim,
+            item['caption']
+        ))
+
+    # 결과 시각화
+    visualize_results(
+        args.query,
+        results,
+        args.output_path,
+        is_image_query
+    )
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision("high")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--query", type=str, required=True, 
+                       help="Path to query image or text query")
+    parser.add_argument("--gallery_json", type=str, required=True,
+                       help="Path to JSON file containing gallery information")
+    parser.add_argument("--output_path", type=str, required=True,
+                       help="Path to save visualization")
+    parser.add_argument("--model", type=str, default='L',
+                       choices=['L-336', 'L', 'B', 'G'])
+    parser.add_argument("--ckpt", type=str, default='',
+                       help="Path to custom checkpoint")
+    parser.add_argument("--new_max_token", type=int, default=248,
+                       help="Maximum number of text tokens")
+    
+    args = parser.parse_args()
+    main(args)