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
|