utils.py

import numpy as np
from torchvision import transforms
import torch
import torch.nn as nn
import torch.nn.functional as F
import PIL
import random
import os
import matplotlib.pyplot as plt
import math
import webdataset as wds

import json
from PIL import Image
import requests
import time
from einops import rearrange

#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def is_interactive():
    import __main__ as main
    return not hasattr(main, '__file__')

def seed_everything(seed=0, cudnn_deterministic=True):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if cudnn_deterministic:
        torch.backends.cudnn.deterministic = True
    else:
        ## needs to be False to use conv3D
        print('Note: not using cudnn.deterministic')

def np_to_Image(x):
    if x.ndim==4:
        x=x[0]
    return PIL.Image.fromarray((x.transpose(1, 2, 0)*127.5+128).clip(0,255).astype('uint8'))

def torch_to_Image(x):
    if x.ndim==4:
        x=x[0]
    return transforms.ToPILImage()(x)

def Image_to_torch(x):
    try:
        x = (transforms.ToTensor()(x)[:3].unsqueeze(0)-.5)/.5
    except:
        x = (transforms.ToTensor()(x[0])[:3].unsqueeze(0)-.5)/.5
    return x

def torch_to_matplotlib(x,device='cpu'):
    if torch.mean(x)>10:
        x = (x.permute(0, 2, 3, 1)).clamp(0, 255).to(torch.uint8)
    else:
        x = (x.permute(0, 2, 3, 1) * 255).clamp(0, 255).to(torch.uint8)
    if device=='cpu':
        return x[0]
    else:
        return x.cpu().numpy()[0]

def batchwise_pearson_correlation(Z, B):
    # Calculate means
    Z_mean = torch.mean(Z, dim=1, keepdim=True)
    B_mean = torch.mean(B, dim=1, keepdim=True)

    # Subtract means
    Z_centered = Z - Z_mean
    B_centered = B - B_mean

    # Calculate Pearson correlation coefficient
    numerator = Z_centered @ B_centered.T
    Z_centered_norm = torch.linalg.norm(Z_centered, dim=1, keepdim=True)
    B_centered_norm = torch.linalg.norm(B_centered, dim=1, keepdim=True)
    denominator = Z_centered_norm @ B_centered_norm.T

    pearson_correlation = (numerator / denominator)
    return pearson_correlation

def batchwise_cosine_similarity(Z,B):
    Z = Z.flatten(1)
    B = B.flatten(1).T
    Z_norm = torch.linalg.norm(Z, dim=1, keepdim=True)  # Size (n, 1).
    B_norm = torch.linalg.norm(B, dim=0, keepdim=True)  # Size (1, b).
    cosine_similarity = ((Z @ B) / (Z_norm @ B_norm)).T
    return cosine_similarity

def prenormed_batchwise_cosine_similarity(Z,B):
    return (Z @ B.T).T

def cosine_similarity(Z,B,l=0):
    Z = nn.functional.normalize(Z, p=2, dim=1)
    B = nn.functional.normalize(B, p=2, dim=1)
    # if l>0, use distribution normalization
    # https://twitter.com/YifeiZhou02/status/1716513495087472880
    Z = Z - l * torch.mean(Z,dim=0)
    B = B - l * torch.mean(B,dim=0)
    cosine_similarity = (Z @ B.T).T
    return cosine_similarity

def topk(similarities,labels,k=5):
    if k > similarities.shape[0]:
        k = similarities.shape[0]
    topsum=0
    for i in range(k):
        topsum += torch.sum(torch.argsort(similarities,axis=1)[:,-(i+1)] == labels)/len(labels)
    return topsum

def get_non_diagonals(a):
    a = torch.triu(a,diagonal=1)+torch.tril(a,diagonal=-1)
    # make diagonals -1
    a=a.fill_diagonal_(-1)
    return a

def gather_features(image_features, voxel_features, accelerator):  
    all_image_features = accelerator.gather(image_features.contiguous())
    if voxel_features is not None:
        all_voxel_features = accelerator.gather(voxel_features.contiguous())
        return all_image_features, all_voxel_features
    return all_image_features



def soft_clip_loss(preds, targs, temp=0.125):
    clip_clip = (targs @ targs.T)/temp
    brain_clip = (preds @ targs.T)/temp
    loss1 = -(brain_clip.log_softmax(-1) * clip_clip.softmax(-1)).sum(-1).mean()
    loss2 = -(brain_clip.T.log_softmax(-1) * clip_clip.softmax(-1)).sum(-1).mean()
    
    loss = (loss1 + loss2)/2
    return loss

def soft_siglip_loss(preds, targs, temp, bias):
    temp = torch.exp(temp)
    
    logits = (preds @ targs.T) * temp + bias
    # diagonals (aka paired samples) should be >0 and off-diagonals <0
    labels = (targs @ targs.T) - 1 + (torch.eye(len(targs)).to(targs.dtype).to(targs.device))

    loss1 = -torch.sum(nn.functional.logsigmoid(logits * labels[:len(preds)])) / len(preds)
    loss2 = -torch.sum(nn.functional.logsigmoid(logits.T * labels[:,:len(preds)])) / len(preds)
    loss = (loss1 + loss2)/2
    return loss

def mixco_hard_siglip_loss(preds, targs, temp, bias, perm, betas):
    temp = torch.exp(temp)
    
    probs = torch.diag(betas)
    probs[torch.arange(preds.shape[0]).to(preds.device), perm] = 1 - betas

    logits = (preds @ targs.T) * temp + bias
    labels = probs * 2 - 1
    #labels = torch.eye(len(targs)).to(targs.dtype).to(targs.device) * 2 - 1
    
    loss1 = -torch.sum(nn.functional.logsigmoid(logits * labels)) / len(preds)
    loss2 = -torch.sum(nn.functional.logsigmoid(logits.T * labels)) / len(preds)
    loss = (loss1 + loss2)/2
    return loss

def mixco(voxels, beta=0.15, s_thresh=0.5, perm=None, betas=None, select=None):
    if perm is None:
        perm = torch.randperm(voxels.shape[0])
    voxels_shuffle = voxels[perm].to(voxels.device,dtype=voxels.dtype)
    if betas is None:
        betas = torch.distributions.Beta(beta, beta).sample([voxels.shape[0]]).to(voxels.device,dtype=voxels.dtype)
    if select is None:
        select = (torch.rand(voxels.shape[0]) <= s_thresh).to(voxels.device)
    betas_shape = [-1] + [1]*(len(voxels.shape)-1)
    voxels[select] = voxels[select] * betas[select].reshape(*betas_shape) + \
        voxels_shuffle[select] * (1 - betas[select]).reshape(*betas_shape)
    betas[~select] = 1
    return voxels, perm, betas, select

def mixco_clip_target(clip_target, perm, select, betas):
    clip_target_shuffle = clip_target[perm]
    clip_target[select] = clip_target[select] * betas[select].reshape(-1, 1) + \
        clip_target_shuffle[select] * (1 - betas[select]).reshape(-1, 1)
    return clip_target



def mixco_nce(preds, targs, temp=0.1, perm=None, betas=None, select=None, distributed=False, 
              accelerator=None, local_rank=None, bidirectional=True):
    brain_clip = (preds @ targs.T)/temp
    
    if perm is not None and betas is not None and select is not None:
        probs = torch.diag(betas)
        probs[torch.arange(preds.shape[0]).to(preds.device), perm] = 1 - betas

        loss = -(brain_clip.log_softmax(-1) * probs).sum(-1).mean()
        if bidirectional:
            loss2 = -(brain_clip.T.log_softmax(-1) * probs.T).sum(-1).mean()
            loss = (loss + loss2)/2
        return loss
    else:
        loss =  F.cross_entropy(brain_clip, torch.arange(brain_clip.shape[0]).to(brain_clip.device))
        if bidirectional:
            loss2 = F.cross_entropy(brain_clip.T, torch.arange(brain_clip.shape[0]).to(brain_clip.device))
            loss = (loss + loss2)/2
        return loss



def temp_consist_loss(masks):
    """
    Args:
        masks (torch.Tensor): Predicted masks of shape (B, T, C, H, W),
                             where B is the batch size, T is the number of frames,
                             C is the number of channels, H and W are the height and width.
    Returns:
        loss (torch.Tensor): Temporal consistency loss.
    """
    masks = rearrange(masks, '(b f) c h w -> b f c h w', f=6)

    B, F, C, H, W = masks.shape
    loss = 0.0

    # Compute L1 loss between consecutive frames
    for t in range(F - 1):
        loss += torch.mean(torch.abs(masks[:, t] - masks[:, t + 1]))

    # Normalize by the number of frame pairs
    return loss / (F - 1)


class DiceLoss(nn.Module):
    def __init__(self, smooth=1e-7):
        super(DiceLoss, self).__init__()
        self.smooth = smooth

    def forward(self, pred, mask):
        """
        pred: [B, 1, H, W]
        mask: [B, 1, H, W]
        """
        assert pred.shape == mask.shape, "pred and mask should have the same shape."
        p = torch.sigmoid(pred)
        intersection = torch.sum(p * mask)
        union = torch.sum(p) + torch.sum(mask)
        dice_loss = (2.0 * intersection + self.smooth) / (union + self.smooth)

        return 1 - dice_loss


def count_params(model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('param counts:\n{:,} total\n{:,} trainable'.format(total, trainable))
    return trainable
    
def check_loss(loss):
    if loss.isnan().any():
        raise ValueError('NaN loss')

def cosine_anneal(start, end, steps):
    return end + (start - end)/2 * (1 + torch.cos(torch.pi*torch.arange(steps)/(steps-1)))

def resize(img, img_size=128):
    if img.ndim == 3: img = img[None]
    return nn.functional.interpolate(img, size=(img_size, img_size), mode='nearest')

pixcorr_preprocess = transforms.Compose([
    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),
])
def pixcorr(images,brains,nan=True):
    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)
    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)
    if nan:
        corrmean = torch.nanmean(torch.diag(batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)))
    else:
        corrmean = torch.mean(torch.diag(batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)))
    return corrmean

def select_annotations(annots, random=True):
    """
    There are 5 annotations per image. Select one of them for each image.
    """
    for i, b in enumerate(annots):
        t = ''
        if random:
            # select random non-empty annotation
            while t == '':
                rand = torch.randint(5, (1,1))[0][0]
                t = b[rand]
        else:
            # select first non-empty annotation
            for j in range(5):
                if b[j] != '':
                    t = b[j]
                    break
        if i == 0:
            txt = np.array(t)
        else:
            txt = np.vstack((txt, t))
    txt = txt.flatten()
    return txt

from generative_models.sgm.util import append_dims
def unclip_recon(x, diffusion_engine, vector_suffix,
                 num_samples=1, offset_noise_level=0.04, device = 'cpu'):
    assert x.ndim==3
    if x.shape[0]==1:
        x = x[[0]]
    with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16), diffusion_engine.ema_scope():
        z = torch.randn(num_samples,4,96,96).to(device) # starting noise, can change to VAE outputs of initial image for img2img

        # clip_img_tokenized = clip_img_embedder(image) 
        # tokens = clip_img_tokenized
        token_shape = x.shape
        tokens = x
        # print(f"\033[92m tokens {tokens.shape} \033[0m")

        c = {"crossattn": tokens.repeat(num_samples,1,1).to(z.device), "vector": vector_suffix.repeat(num_samples,1).to(z.device)}

        tokens = torch.randn_like(x)
        uc = {"crossattn": tokens.repeat(num_samples,1,1).to(z.device), "vector": vector_suffix.repeat(num_samples,1).to(z.device)}

        for k in c:
            c[k], uc[k] = map(lambda y: y[k][:num_samples].to(device), (c, uc))

        noise = torch.randn_like(z)
        sigmas = diffusion_engine.sampler.discretization(diffusion_engine.sampler.num_steps)
        sigma = sigmas[0].to(z.device)

        if offset_noise_level > 0.0:
            noise = noise + offset_noise_level * append_dims(
                torch.randn(z.shape[0], device=z.device), z.ndim
            )
        noised_z = z + noise * append_dims(sigma, z.ndim)
        noised_z = noised_z / torch.sqrt(
            1.0 + sigmas[0] ** 2.0
        )  # Note: hardcoded to DDPM-like scaling. need to generalize later.

        def denoiser(x, sigma, c):
            return diffusion_engine.denoiser(diffusion_engine.model, x, sigma, c)

        samples_z = diffusion_engine.sampler(denoiser, noised_z, cond=c, uc=uc)

        # print(f"\033[92m samples_z {samples_z.shape} \033[0m")
        samples_x = diffusion_engine.decode_first_stage(samples_z)

        # print(f"\033[92m samples_x {samples_x.shape} \033[0m")


        # samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
        samples = torch.clamp((samples_x*.8+.2), min=0.0, max=1.0)
        return samples

#  Numpy Utility 
def iterate_range(start, length, batchsize):
    batch_count = int(length // batchsize )
    residual = int(length % batchsize)
    for i in range(batch_count):
        yield range(start+i*batchsize, start+(i+1)*batchsize),batchsize
    if(residual>0):
        yield range(start+batch_count*batchsize,start+length),residual 
        
# Torch fwRF
def get_value(_x):
    return np.copy(_x.data.cpu().numpy())

def soft_cont_loss(student_preds, teacher_preds, teacher_aug_preds, temp=0.125):
    teacher_teacher_aug = (teacher_preds @ teacher_aug_preds.T)/temp
    teacher_teacher_aug_t = (teacher_aug_preds @ teacher_preds.T)/temp
    student_teacher_aug = (student_preds @ teacher_aug_preds.T)/temp
    student_teacher_aug_t = (teacher_aug_preds @ student_preds.T)/temp

    loss1 = -(student_teacher_aug.log_softmax(-1) * teacher_teacher_aug.softmax(-1)).sum(-1).mean()
    loss2 = -(student_teacher_aug_t.log_softmax(-1) * teacher_teacher_aug_t.softmax(-1)).sum(-1).mean()
    
    loss = (loss1 + loss2)/2
    return loss


# ============================================================================
# 类别字典 (Category Dictionaries)
# 支持 CC2017 (51类)、NSD/COCO (80类)、HCP (577类)
# ============================================================================

# CC2017: 51 类 (0-indexed)
CC2017_CLS_DICT = {
    0: 'animal', 1: 'human', 2: 'vehicle', 3: 'building', 4: 'clothing',
    5: 'weapon', 6: 'plant', 7: 'appliance', 8: 'tool', 9: 'container',
    10: 'body part', 11: 'furniture', 12: 'device', 13: 'fabric', 14: 'fruit',
    15: 'vegetable', 16: 'insect', 17: 'landscape feature', 18: 'water body',
    19: 'organism', 20: 'fish', 21: 'reptile', 22: 'mammal', 23: 'accessory',
    24: 'sports equipment', 25: 'food', 26: 'drink', 27: 'light source',
    28: 'weather phenomenon', 29: 'jewelry', 30: 'musical instrument',
    31: 'structure', 32: 'flying vehicle', 33: 'toy', 34: 'kitchen item',
    35: 'writing tool', 36: 'gardening tool', 37: 'scientific equipment',
    38: 'furniture accessory', 39: 'roadway', 40: 'weaponry accessory',
    41: 'sports field', 42: 'money', 43: 'timekeeping device', 44: 'decoration',
    45: 'art', 46: 'stationery', 47: 'kitchen appliance', 48: 'rock/mineral',
    49: 'soil/substrate', 50: 'climate/atmosphere component'
}

# NSD/COCO: 80 类 (1-indexed)
NSD_CLS_DICT = {
    1: 'person', 2: 'bicycle', 3: 'car', 4: 'motorcycle', 5: 'airplane',
    6: 'bus', 7: 'train', 8: 'truck', 9: 'boat', 10: 'traffic light',
    11: 'fire hydrant', 13: 'stop sign', 14: 'parking meter', 15: 'bench',
    16: 'bird', 17: 'cat', 18: 'dog', 19: 'horse', 20: 'sheep',
    21: 'cow', 22: 'elephant', 23: 'bear', 24: 'zebra', 25: 'giraffe',
    27: 'backpack', 28: 'umbrella', 31: 'handbag', 32: 'tie', 33: 'suitcase',
    34: 'frisbee', 35: 'skis', 36: 'snowboard', 37: 'sports ball', 38: 'kite',
    39: 'baseball bat', 40: 'baseball glove', 41: 'skateboard', 42: 'surfboard',
    43: 'tennis racket', 44: 'bottle', 46: 'wine glass', 47: 'cup', 48: 'fork',
    49: 'knife', 50: 'spoon', 51: 'bowl', 52: 'banana', 53: 'apple',
    54: 'sandwich', 55: 'orange', 56: 'broccoli', 57: 'carrot', 58: 'hot dog',
    59: 'pizza', 60: 'donut', 61: 'cake', 62: 'chair', 63: 'couch',
    64: 'potted plant', 65: 'bed', 67: 'dining table', 70: 'toilet',
    72: 'tv', 73: 'laptop', 74: 'mouse', 75: 'remote', 76: 'keyboard',
    77: 'cell phone', 78: 'microwave', 79: 'oven', 80: 'toaster', 81: 'sink',
    82: 'refrigerator', 84: 'book', 85: 'clock', 86: 'vase', 87: 'scissors',
    88: 'teddy bear', 89: 'hair drier', 90: 'toothbrush'
}

# HCP: 577 类 (0-indexed)
HCP_CLS_DICT = {i: cat for i, cat in enumerate([
    'trunk', 'hoodie', 'railroad track', 'beach', 'bone', 'tag', 'water', 'baseball player', 'toy',
    'thong', 'angelfish', 'helm', 'trailer', 'owl', 'man', 'clothing', 'rack', 'bag', 'pearl', 'tank',
    'blazer', 'ground', 'tugboat', 'tie', 'fire ball', 'skier', 'excavator', 'bike', 'cage', 'pad',
    'fingernail', 'center', 'grill', 'fire', 'hat', 'baby', 'landing gear', 'tunnel', 'eel', 'bale',
    'sand', 'lobby', 'student', 'bison', 'mouse', 'strawberry', 'chihuahua', 'hamburger', 'tusk', 'nose',
    'tug boat', 'necklace', 'crowd', 'tv', 'glasses', 'goat', 'car', 'fork', 'adult', 'blanket', 'rv',
    'subway', 'construction site', 'card', 'surfer', 'arm', 'redwood', 'cloth', 'fall', 'sun', 'mouth',
    'golf ball', 'event', 'ski', 'ball', 'edge', 'garage', 'leaf', 'cake', 'tank top', 'pirate', 'airport',
    'donut', 'kick', 'plate', 'step', 'pose', 'wrist', 'house', 'door', 'wire', 'cap', 'hair', 'boulder',
    'bridge', 'pant', 'scissor', 'court', 'phone', 'soccer ball', 'belt', 'color', 'art', 'sign', 'sailboat',
    'headphone', 'woman', 'salad', 'bird', 'coffee cup', 'background', 'basketball', 'fighter jet', 'wheelchair',
    'railroad', 'tray', 'cruise ship', 'airplane', 'substance', 'position', 'cheek', 'slice', 'fireplace',
    'column', 'parking lot', 'paintbrush', 'whale shark', 'stick', 'mask', 'field', 'cupcake', 'truck', 'dock',
    'dirt', 'mirror', 'tarp', 'bowl', 'firework', 'test tube', 'face', 'leather', 'dolphin', 'horse', 'bear',
    'desk', 'back', 'hospital', 'printer', 'spoon', 'exercise', 'character', 'frisbee', 'screen', 'doorway',
    'tennis ball', 'combine', 'tree', 'pencil', 'palm tree', 'microphone', 'valley', 'bra', 'drawing', 'meal',
    'rock', 'deer', 'hill', 'coral reef', 'jetliner', 'brain', 'microscope', 'street', 'server rack', 'stocking',
    'harvester', 'skyline', 'jacket', 'ship', 'motion', 'cell phone', 'knee', 'branch', 'bed', 'floor', 'swamp',
    'lab coat', 'stack', 'air', 'brush', 'astronaut', 'locker', 'anemone', 'town', 'reef', 'overall', 'push up',
    'collar', 'breakfast', 'hand', 'cherry blossom', 'roller coaster', 'road', 'light', 'mustache', 'rubble',
    'dreadlock', 'shirt', 'steering wheel', 'coral', 'cow', 'weed', 'apple', 'gymnasium', 'child', 'bull', 'bun',
    'bow', 'panda cub', 'instrument', 'leg', 'base', 'city', 'revolver', 'legging', 'doll', 'cart', 'vegetable',
    'fountain', 'pen', 'dog', 'jean', 'game', 'sheep', 'paw', 'whale', 'sunflower', 'shuttle', 'golf', 'lipstick',
    'seaweed', 'ice', 'antelope', 'living room', 'writing', 'hole', 'forest', 'tablet', 'paint', 'gazelle',
    'bouquet', 'mangrove', 'camera', 'harbor', 'panda', 'path', 'rapid', 'train', 'compass', 'sky', 'woods',
    'alga', 'turkey', 'box', 'number', 'dumbbell', 'stingray', 'surface', 'statue', 'vest', 'room', 'umbrella',
    'balloon', 'precinct', 'nurse', 'fish', 'bat', 'beak', 'shark', 'duck', 'short', 'head', 'painting', 'karate',
    'ice cream', 'vehicle', 'labrador retriever', 'speedometer', 'earth', 'guitar', 'machine', 'lake', 'dinner',
    'pillar', 'blossom', 'shotgun', 'eye', 'range', 'world', 'golf course', 'shack', 'night', 'aquarium',
    'counter top', 'couch', 'platter', 'bamboo', 'rose', 'liquid', 'iguana', 't - shirt', 'scuba diver', 'roof',
    'shower', 'cattle', 'darkness', 'text', 'boat', 'tongue', 'fishing rod', 'juice', 'ocean', 'sweater', 'gun',
    'mud', 'stable', 'cloud', 'shoulder', 'people', 'chicken', 'baseball field', 'melon', 'wing', 'hillside',
    'paper', 'candle', 'housewife', 'wood', 'sling', 'apron', 'chair', 'station', 'finger', 'flame', 'comforter',
    'push', 'dark', 'croissant', 'puppy', 'body', 'bridle', 'coffee maker', 'clock', 'plant', 'site', 'grass',
    'skirt', 'syringe', 'egg', 'something', 'dresser', 'liberty', 'cliff', 'counter', 'suv', 'neck', 'market',
    'towel', 'island', 'cigarette', 'sweatshirt', 'backpack', 'bush', 'yard', 'panel', 'hut', 'monkey', 'zoo',
    'book', 'toilet', 'lip', 'butterfly', 'tail', 'elephant', 'shop', 'dress', 'jeep', 'explosion', 'beard',
    'word', 'river', 'shelf', 'rifle', 'shore', 'stream', 'stomach', 'area', 'motorcycle', 'row', 'badge',
    'mother', 'stand', 'piano', 'luggage', 'cup', 'track', 'hallway', 'wine', 'crane', 'television', 'sponge',
    'wind', 'mountain', 'parrot fish', 'golf club', 'bill', 'girl', 'wheat', 'suspender', 'rod', 'sink', 'family',
    'snow', 'pastry', 'hallows', 'stair', 'person', 'player', 'arrow', 'animal', 'tooth', 'father', 'table',
    'slope', 'space station', 'parrotfish', 'combine harvester', 'boy', 'elevator', 'hammock', 'video', 'outfit',
    'bench', 'top', 'bunker', 'bottle', 'server', 'rug', 'grain', 'tower', 'image', 'van', 'buffalo', 'restaurant',
    'banana', 'soccer', 'jellyfish', 'penguin', 'heel', 'knife', 'gear', 'coffee', 'gravel', 'ice cream cone',
    'smile', 'park', 'object', 'dollar', 'smoke', 'suitcase', 'flower', 'highway', 'tent', 'beanie', 'orange',
    'pile', 'glove', 'computer', 'dough', 'cave', 'fruit', 'oven', 'ipad', 'tuxedo', 'laptop', 'footage',
    'basketball hoop', 'scarf', 'lettuce', 'moon', 'board', 'present', 'piece', 'meat', 'caribbean', 'shoe',
    'fern', 'mango', 'platform', 'storm', 'pilgrim', 'bottom', 'beer', 'office', 'quote', 'sidewalk', 'desert',
    'hay', 'lion', 'crate', 'figurine', 'classroom', 'trash', 'panda bear', 'arch', 'barn', 'harness', 'flag',
    'club', 'turtle', 'straw', 'marsh', 'wall', 'None', 'hay bale', 'building', 'window', 'suit', 'blueberry',
    'rain', 'walkway', 'wheel', 'spacesuit', 'bread', 'mat', 'gym', 'village', 'waterfall', 'glass', 'tractor',
    'bandage', 'data center', 'carriage', 'kitchen', 'bar', 'niagara falls', 'skateboard', 'butterfly fish',
    'bookshelf', 'breast', 'bow tie', 'coat', 'wave', 'fence', 'trail', 'stair case', 'line', 'diver', 'food',
    'canada', 'biker', 'hoop', 'espresso machine', 'cat', 'money', 'foot', 'snowboard', 'ant', 'pond', 'uniform',
    'lynx', 'satellite', 'robe', 'ad', 'sunglasses', 'pizza', 'cobblestone', ''
])}


def get_cls_dict(dataset):
    """根据数据集名称返回对应的类别字典"""
    if dataset == 'cc2017':
        return CC2017_CLS_DICT
    elif dataset == 'nsd':
        return NSD_CLS_DICT
    elif dataset == 'hcp':
        return HCP_CLS_DICT
    else:
        raise ValueError(f"Unknown dataset: {dataset}")


def get_num_classes(dataset):
    """返回数据集的类别数"""
    if dataset == 'cc2017':
        return 51
    elif dataset == 'nsd':
        return 80
    elif dataset == 'hcp':
        return 577
    else:
        raise ValueError(f"Unknown dataset: {dataset}")