jittor-FSC/JCLIP/utils.py

import random
import yaml
import jittor as jt
import numpy as np
import pandas as pd
from PIL import Image
import jclip as clip
from jittor.transform import ImageNormalize, Compose
from datetime import datetime
from jittor.optim import Optimizer

# 类别到idx的映射
def class_2_idx(path):
    '''
    path   : classes.txt文件路径
    return : dict: {'Bear': 0,'Bee': 1, ..., 'papillon': 373,...}
    '''
    df_classes   = pd.read_csv(path, delimiter = ' ', header = None, index_col = False)
    list_classes = list(map(lambda x: x.split('_',1)[-1], df_classes[0]))
    res = {k:v for k,v in zip(list_classes,range(len(list_classes)))}
    
    return res


def generate_prompt(name :str):
    '''
    用于生成提示语句
    '''
    li = [name,
        f"{name} , {len(name)}",
        f"A photo of a {name}",
        f"A photo of a {name} with {len(name)}",
        f"A photo of {name} and the length of the prompt is {len(name)}",
        f"A photo of {name} and the length of the {name} is {len(name)}",
        f"A photo of {name.replace('_', ' ')} and the length of the name is {len(name.replace('_', ' '))}",
        f"A photo of {name} with {len(name)} and _ in {name}" if '_' in name else f"A photo of {name} with {len(name)} and _ not in {name}",
        f"A photo of {name}, {name}, {name}"]


    return li[4]


def normalize_tensor(tensor, a=0.0, b=1.0):
    '''
    归一化
    '''
    tensor_min = tensor.min()
    tensor_max = tensor.max()
    normalized_tensor = a + (tensor - tensor_min) * (b - a) / (tensor_max - tensor_min)
    
    return normalized_tensor


def get_val_text_features(classes_path, model):
    '''
    用于获取模型在待预测类别上的文本特征
    '''
    df_classes     = pd.read_csv(classes_path , delimiter = ' ' , header = None , index_col = False)
    classes        = list(map(lambda x: generate_prompt(x.split('_',1)[-1]), df_classes[0]))
    text           = clip.tokenize(classes)
    with jt.no_grad():
        text_features  = model.encode_text(text)
        text_features /= text_features.norm(dim=-1, keepdim=True)     # shape -> (374, 512)  包含了所有类别的文本特征
    return text_features                             


# 自定义的随机擦除函数：类似PyTorch中的torchvision.transforms.RandomErasing函数
class RandomErasing:
    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0):
        self.p = p
        self.scale = scale
        self.ratio = ratio
        self.value = value

    def __call__(self, img):
        if random.uniform(0, 1) > self.p:
            return img

        img = np.array(img)
        h, w, c = img.shape
        area = h * w

        for attempt in range(100):
            target_area = random.uniform(*self.scale) * area
            aspect_ratio = random.uniform(*self.ratio)

            erase_h = int(round(np.sqrt(target_area * aspect_ratio)))
            erase_w = int(round(np.sqrt(target_area / aspect_ratio)))

            if erase_h < h and erase_w < w:
                x1 = random.randint(0, h - erase_h)
                y1 = random.randint(0, w - erase_w)

                if self.value == 'random':
                    img[x1:x1+erase_h, y1:y1+erase_w, :] = np.random.randint(0, 256, (erase_h, erase_w, c), dtype=np.uint8)
                else:
                    img[x1:x1+erase_h, y1:y1+erase_w, :] = self.value

                return jt.array(img)

        return jt.array(img)





# 计算模型的参数量 : 根据赛题要求不能超过 500Mb 
def count_parameters_in_mb(model):
    total_params = 0
    for param in model.parameters():
        param_count = param.numel()
        total_params += param_count

    total_params_in_mb = total_params / 1e6  # 转换为百万参数（Mb）
    
    assert total_params_in_mb <= 500, "Model {} has too many parameters: {:.2f} Mb. The limit is 500 Mb.".format(model.__class__.__name__, total_params_in_mb)
    # print("Model {} : {:.2f} Mb".format(model.__class__.__name__, total_params_in_mb))
    
    return total_params_in_mb



# 格式化输出args对象中除了 exclude_keys 以外的键值对
def format_args(args, exclude_keys=['root_TrainSet', 'train_path', 'save_path', 'TestSetZ_path', 'label_path', 'classes_path', 'classes_b_path', 'class_4_path']):
    
    max_key_length = max(len(key) for key in vars(args).keys())     # 计算所有键的最大长度
    formatted_args = "\n".join([f"{key.ljust(max_key_length)} : {value}" for key, value in vars(args).items() if key not in exclude_keys]) # 格式化键值对
    
    return formatted_args





# 重写clip.py文件中【_transform方法】，,主要更新了对图像的裁剪的过程
class Image_Transform():
    def __init__(self, img:Image, clip_init):
        self.img  = clip.clip.Resize(224, mode=Image.BICUBIC)(img)  
        self.size = self.img.size
        self.clip = clip_init 
        
    def choose_best_img(self):
    
        # 若图像的长宽比在指定的阈值以内,则直接缩放;否则就裁剪多张图片送入CLIP中，选取预测类别的众数中预测概率值最大的那张
        # 比如同一张图片裁剪了10张图片，预测到了8个A，1个B和1个C，那么选择8个A中概率值最大的那张
        if self.is_direct_scaling(self.size):
            img = self.img.resize((224,224))
            img = self.transform()(img)
            return img
        
        else:
            self.class_path= r'F:\jittor_comprtition\Competition1\classes.txt'
            classes        = list(map(lambda x:generate_prompt(x.split('_',1)[-1]),list(class_2_idx(self.class_path).keys())))
            text           = clip.tokenize(classes)
            with jt.no_grad():
                text_features  = self.clip.encode_text(text)
                text_features /= text_features.norm(dim=1, keepdim=True)
                self.text_features = text_features

            crop_list = self.crop_image_sliding(self.img, 10)

            img_tensor = []
            for img in crop_list:
                img_tensor.append(self.transform()(img))
            img_tensor = jt.array(img_tensor)    # [crop_num, 3, 224 224]
            with jt.no_grad():
                image_features    = self.clip.encode_image(img_tensor)
                image_features   /= image_features.norm(dim=1, keepdim=True)
                text_probs = (100.0 * image_features @ self.text_features.transpose(0, 1)).softmax(dim=-1)
                pred_probs, top_labels = text_probs.topk(1)

                # 选取预测类别的众数且预测概率值最大的那张图片,如果有多个众数，则随机选择一个
                idx, _, idx_counts = jt.unique(top_labels, return_inverse=True, return_counts=True) # 统计众数及其数量
                idx_mode           = idx[jt.argmax(idx_counts, dim=0)[0]]   
                index_mode         = jt.equal(idx_mode, top_labels.flatten()).nonzero().flatten()
                best_img_index     = index_mode[jt.argmax(pred_probs[index_mode],dim=0)[0]].item()
                best_img           = img_tensor[best_img_index]
                return best_img



    def is_direct_scaling(self, img_size:tuple, threshold=1.05):
        aspect_ratio = max(img_size) / 224.0
        return aspect_ratio <= threshold

    def transform(self):
        return Compose([clip.clip._convert_image_to_rgb,
            ImageNormalize((0.48145466, 0.4578275, 0.40821073),
                        (0.26862954, 0.26130258, 0.27577711)),
            clip.clip.ImageToTensor()
        ])
    
    def predict(self, img_fea):
        text_probs = (100.0 * img_fea @ self.text_features.transpose(0, 1)).softmax(dim=-1)

    
    def crop_image_sliding(self, img, crop_num, step=None):

        width, height = img.size
        
        # 检查维度，确定沿哪个维度滑动裁剪
        if width > 224:
            slide_dim = 'width'
            fixed_dim = 224
        else:
            slide_dim = 'height'
            fixed_dim = 224

        
        # 计算默认步长，如果没有指定
        if step is None:
            if slide_dim == 'width':
                step = (width - 224) // (crop_num - 1)
            else:
                step = (height - 224) // (crop_num - 1)
        
        cropped_images = []
        
        for i in range(crop_num):
            if slide_dim == 'width':
                left = i * step
                if left + 224 > width:
                    left = width - 224
                box = (left, 0, left + 224, 224)
            else:
                top = i * step
                if top + 224 > height:
                    top = height - 224
                box = (0, top, 224, top + 224)
            
            cropped_img = img.crop(box)
            cropped_images.append(cropped_img)
        
        return cropped_images


# 获取当天日期
def get_date_format():
    today = datetime.now()
    formatted_date =  str(today.month).zfill(2) + str(today.day).zfill(2)
    return formatted_date


# 从yaml文件中读取参数
def load_yaml_params(yaml_path, args):
    with open(yaml_path, 'rb') as f:
        config = yaml.safe_load(f)
    for key in config.keys():
        assert hasattr(args, key), f"The 'args' object does not have the attribute: {key}"
        setattr(args, key, config[key])
    return args