from torchvision import transforms
from utils import custom_transforms as tr
import numpy as np
import pandas as pd


def transform_ts(args,sample):
    #将图像从cv读的格式转为归一化，并转为tensor
    composed_transforms = transforms.Compose([
        tr.FixedResize(size=args['crop_size']),
        tr.Normalize(mean=(0.335, 0.358, 0.332), std=(0.141, 0.138, 0.143)),
        tr.ToTensor()])
    return composed_transforms(sample)



def colour_code_segmentation(image, label_values):
    label_values = [label_values[key] for key in label_values]
    colour_codes = np.array(label_values)
    x = colour_codes[image.astype(int)]
    return x


def get_label_info(csv_path):
    ann = pd.read_csv(csv_path)
    label = {}
    for iter, row in ann.iterrows():
        label_name = row['name']
        r = row['r']
        g = row['g']
        b = row['b']
        label[label_name] = [int(r), int(g), int(b)]
    return label