34 lines
943 B
Python
34 lines
943 B
Python
|
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
|