def train(self):
'''
训练模型，必须实现此方法
:return:
'''
dataset = CarDataset()
# 24 + 1
model = fasterRCNN(num_classes=25)
load_pretrained_weights(model, './weights/resnet50-19c8e357.pth')
model = model.cuda()
model.train()
# fang[-1]
optimizer = build_optimizer(model, optim='adam')
max_epoch = args.EPOCHS
batch_size = args.BATCH
scheduler = build_scheduler(optimizer,
lr_scheduler='cosine',
max_epoch=max_epoch)
train_loader = DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=self.my_collate)
cudnn.benchmark = True
for epoch in range(max_epoch):
for index, data in enumerate(train_loader):
ims, cls_labels, bbox_labels = data
targets = []
for i in range(len(ims)):
d = {}
d['labels'] = torch.tensor(cls_labels[i],
dtype=torch.long).cuda()
d['boxes'] = torch.tensor(bbox_labels[i],
dtype=torch.long).cuda()
targets.append(d)
ims = torch.tensor([im.cpu().detach().numpy() for im in ims])
ims = ims.cuda()
out = model(ims, targets)
loss_classifier = out['loss_classifier']
loss_box_reg = out['loss_box_reg']
loss_objectness = out['loss_objectness']
loss_rpn_box_reg = out['loss_rpn_box_reg']
loss = 0.5 * loss_classifier + 5 * loss_box_reg + loss_objectness + 10 * loss_rpn_box_reg
loss.backward()
optimizer.step()
if index % 10 == 0:
print(
"Epoch: [{}/{}][{}/{}] Loss: loss_classifier: {:.2f}, "
"loss_box_reg: {:.2f}, loss_objectness: {:.2f}, "
"loss_rpn_box_reg: {:.2f}, total loss: {:.2f}".format(
epoch + 1, max_epoch, index + 1, len(train_loader),
loss_classifier, loss_box_reg, loss_objectness,
loss_rpn_box_reg, loss))
# n_iter = epoch*len(train_loader) + index
# writer.add_scalar('loss', loss, n_iter)
scheduler.step()
if (epoch + 1) % 100 == 0:
torch.save(model.state_dict(), 'last.pth')
def predict(self, image_path):
'''
模型预测返回结果
:参数示例 image_path='./data/input/image/0.jpg'
:return: 返回预测结果格式具体说明如下：
'''
# 0.jpg
image_name = os.path.basename(image_path)
model = self.load_model()
model = model.cuda()
model.eval()
dataset = CarDataset()
# test_dataloader = DataLoader(dataset, batch_size=1, shuffle=False)
pred_result = []
# for index, data in enumerate(test_dataloader):
# im, cls_label, gt_bbox = data
# im = im.cuda()
# out = model(im, None)[0]
# bboxes_pred = out['boxes']
# labels_pred = out['labels']
# scores_pred = out['scores']
# bboxes_pred = bboxes_pred.tolist()
# # bboxes_pred = [bbox.tolist() for bbox in bboxes_pred]
# for index, bbox in enumerate(bboxes_pred):
# d = {}
# d["image_name"] = image_name
# label = labels_pred[index].cpu().detach().numpy()
# # print(label)
# # print(dataset.greek_nums_index_map)
# label_name = dataset.greek_nums_index_map[str(label)]
# d['label_name'] = label_name
# xmin, ymin, xmax, ymax = bbox[:]
# d['bbox'] = [int(xmin), int(ymin), int(xmax-xmin+1), int(ymax-ymin+1)]
# score = scores_pred[index].cpu().detach().numpy()
# d['confidence'] = score
# pred_result.append(d)
im = Image.open(image_path)
im = dataset.transforms(im)
im = im.unsqueeze(0)
im = im.cuda()
out = model(im, None)[0]
bboxes_pred = out['boxes']
labels_pred = out['labels']
scores_pred = out['scores']
bboxes_pred = bboxes_pred.tolist()
for index, bbox in enumerate(bboxes_pred):
d = {}
d["image_name"] = image_name
label = labels_pred[index].cpu().detach().numpy()
label_name = dataset.greek_nums_index_map[str(label)]
d['label_name'] = label_name
xmin, ymin, xmax, ymax = bbox[:]
d['bbox'] = [
int(xmin),
int(ymin),
int(xmax - xmin + 1),
int(ymax - ymin + 1)
]
score = scores_pred[index].cpu().detach().numpy().tolist()
d['confidence'] = score
pred_result.append(d)
# 返回bbox格式为 [xmin, ymin, width, height]
# pred_result = [{"image_name": image_name, "label_name": 'I', "bbox": [735, 923, 35, 75], "confidence": 0.2},
# {"image_name": image_name, "label_name": 'I', "bbox": [525, 535, 53, 54], "confidence": 0.3}]
# print(pred_result)
return pred_result
from torch.utils.data import DataLoader
import pandas as pd
from tqdm import tqdm
from config import Config
from dataset import CarDataset
from model import CentResnet
from utils import extract_coords, coords2str
PATH = Config.PATH
device = Config.device
predictions = []
test_images_dir = PATH + 'test_images/{}.jpg'
test = pd.read_csv(PATH + 'sample_submission.csv')
df_test = test
test_dataset = CarDataset(df_test, test_images_dir, False)
test_loader = DataLoader(dataset=test_dataset,
batch_size=4,
shuffle=False,
num_workers=4)
model = CentResnet(8).to(device)
model.load_state_dict(torch.load(Config.save_path))
model.eval()
if __name__ == '__main__':
for img, _, _ in tqdm(test_loader):
with torch.no_grad():
output = model(img.to(device))
output = output.data.cpu().numpy()
for out in output:
coords = extract_coords(out)
def train_net(net,
device,
epochs=5,
batch_size=1,
lr=0.001,
val_percent=0.1,
save_cp=True):
dataset = CarDataset.create(augment_transform, mask_transform)
n_val = int(len(dataset) * val_percent)
n_train = len(dataset) - n_val
train, val = random_split(dataset, [n_train, n_val])
train_loader = DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
val_loader = DataLoader(val, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True, drop_last=True)
writer = SummaryWriter(comment=f'LR_{lr}_BS_{batch_size}')
global_step = 0
logging.info(f'''Starting training:
Epochs: {epochs}
Batch size: {batch_size}
Learning rate: {lr}
Training size: {n_train}
Validation size: {n_val}
Checkpoints: {save_cp}
Device: {device.type}
''')
optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min' if net.n_classes > 1 else 'max', patience=2)
if net.n_classes > 1:
criterion = nn.CrossEntropyLoss()
else:
criterion = nn.BCEWithLogitsLoss()
for epoch in range(epochs):
net.train()
epoch_loss = 0
with tqdm(total=n_train, desc=f'Epoch {epoch + 1}/{epochs}', unit='img') as pbar:
for batch in train_loader:
imgs = batch[0]
true_masks = batch[1]
assert imgs.shape[1] == net.n_channels, \
f'Network has been defined with {net.n_channels} input channels, ' \
f'but loaded images have {imgs.shape[1]} channels. Please check that ' \
'the images are loaded correctly.'
imgs = imgs.to(device=device, dtype=torch.float32)
mask_type = torch.float32 if net.n_classes == 1 else torch.long
true_masks = true_masks.to(device=device, dtype=mask_type)
masks_pred = net(imgs)
loss = criterion(masks_pred, true_masks)
epoch_loss += loss.item()
writer.add_scalar('Loss/train', loss.item(), global_step)
pbar.set_postfix(**{'loss (batch)': loss.item()})
optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_value_(net.parameters(), 0.1)
optimizer.step()
pbar.update(imgs.shape[0])
global_step += 1
try:
os.mkdir(dir_checkpoint)
logging.info('Created checkpoint directory')
except OSError:
pass
torch.save(net.state_dict(),
dir_checkpoint + f'CP_epoch{epoch + 1}.pth')
logging.info(f'Checkpoint {epoch + 1} saved !')
writer.close()
import matplotlib.pyplot as plt
import pandas as pd
import torch
from sklearn.model_selection import train_test_split
from config import Config
from dataset import CarDataset
import numpy as np
from model import CentResnet
PATH = Config.PATH
train = pd.read_csv(PATH + 'train.csv')
train_images_dir = PATH + 'train_images/{}.jpg'
df_train, df_dev = train_test_split(train, test_size=0.01, random_state=42)
dev_dataset = CarDataset(df_dev, train_images_dir)
device = Config.device
model = CentResnet(8).to(device)
model.load_state_dict(torch.load(Config.save_path))
img, mask, regr = dev_dataset[0]
plt.figure(figsize=(16, 16))
plt.title('Input image')
plt.imshow(np.rollaxis(img, 0, 3))
plt.show()
plt.figure(figsize=(16, 16))
plt.title('Ground truth mask')
plt.imshow(mask)
if not os.path.isdir(output_path):
os.makedirs(output_path)
image_base_name = image_name.split('.')[0]
new_image_name = image_base_name + '.png'
outfile = os.path.join(output_path, new_image_name)
create_thumbnail_image(img, outfile, size=size)
"""
## CAR196
root_dir = "/data1/Guoxian_Dai/car196"
image_txt = "/data1/Guoxian_Dai/car196/cars_annos.mat"
train_test_split_txt = "/data1/Guoxian_Dai/CUB_200_2011/train_test_split.txt"
label_txt = "/data1/Guoxian_Dai/CUB_200_2011/image_class_labels.txt"
offset = 1
car_dataset = CarDataset(root_dir, image_txt)
train_img_list = car_dataset.train_img_list
train_label_list = car_dataset.train_label_list
test_img_list = car_dataset.test_img_list
test_label_list = car_dataset.test_label_list
output_dir = './car196_thumbnail'
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
"""
print(len(train_img_list))
print(train_img_list[-2:])
print(len(train_label_list))
print(train_label_list[-2:])
print(len(test_img_list))
print(test_img_list[-2:])
def execute(device, model_name, dataset_dir, batch_size, epochs, learning_rate, num_classes, feature_extract, use_pretrained, save_file, print_plots):
# vehicles classes
classes = ('Full size car', 'Mid size car', 'Cross over', 'Van',
'Coupe', 'Family car', 'Beetle', 'Single seater', 'City car', 'Pick up')
model, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained)
model = model.to(device)
print('-' * 100)
print(f"MODEL ARCHITECTURE [{model_name}] [feature extract: {feature_extract}, use_pretrained: {use_pretrained} learning rate: {learning_rate}]")
# print('-' * 100)
# print(model)
# params
params_to_update = model.parameters()
print('-' * 100)
print("PARAMS TO LEARN")
print('-' * 100)
if feature_extract:
params_to_update = []
for name, param in model.named_parameters():
if param.requires_grad:
params_to_update.append(param)
print("\t", name)
else:
for name, param in model.named_parameters():
if param.requires_grad:
print("\t", name)
transform_dict = {
'train': transforms.Compose(
[
transforms.Resize(input_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'eval': transforms.Compose(
[
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
}
# dataset
print('-' * 100)
print(f"IMPORTING DATASET...")
print('-' * 100)
trainset = CarDataset(dataset_dir=dataset_dir / 'train',
transform=transform_dict['train'])
trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, drop_last=True)
testset = CarDataset(dataset_dir=dataset_dir / 'test',
transform=transform_dict['eval'])
testloader = DataLoader(testset, batch_size=batch_size, shuffle=True, drop_last=True)
print('-' * 100)
print(f"DATASET IMPORTED")
print('-' * 100)
# statistics
train_losses = []
train_accs = []
test_losses = []
test_accs = []
true_list = []
predict_list = []
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(params_to_update, lr=learning_rate)
d = datetime.datetime.now()
date_id = str(getattr(d, 'month')) + '-' + str(getattr(d, 'day')) + '__' + str(
getattr(d, 'hour')) + ':' + str(getattr(d, 'minute'))
id = model_name + '__' + date_id + '__' + str(learning_rate) + 'e__'
if feature_extract:
id += 'FE__'
best_test_acc = 0.
for epoch in range(1, epochs):
print(f"EPOCH {epoch}/{epochs}")
# training
print("Training")
train_loss, train_acc = train(trainloader, model, batch_size, criterion, optimizer, device)
train_losses.append(train_loss)
train_accs.append(train_acc)
print(f"EPOCH {epoch + 1}: [TRAINING loss: {train_loss:.5f} acc: {train_acc:.2f}%]")
# test
print("Testing")
with torch.no_grad():
test_loss, test_acc, true, predict = test(testloader, model, batch_size, criterion, device)
if test_acc > best_test_acc:
best_test_acc = test_acc
# Keep track of best epoch
best_epoch_file = Path('data') / (id + 'best_epoch.txt')
with best_epoch_file.open('wt') as f:
f.write(str(epoch))
# Save best model weights
torch.save(model.state_dict(), 'data/' + id + 'weights.pth')
conf_matrix1(id, true, predict, classes, list(i for i in range(10)), figsize=(10, 10))
test_losses.append(test_loss)
test_accs.append(test_acc)
true_list += true
predict_list += predict
print(f"EPOCH {epoch}: [TESTING loss: {test_loss:.5f} acc: {test_acc:.2f}%]")
print('-' * 100)
if save_file:
main.write('data/' + id + 'train_loss.csv', train_losses)
main.write('data/' + id + 'train_acc.csv', train_accs)
main.write('data/' + id + 'test_loss.csv', test_losses)
main.write('data/' + id + 'test_acc.csv', test_accs)
main.write('data/' + id + 'true.csv', true_list)
main.write('data/' + id + 'predict.csv', predict_list)
if print_plots:
plots.printPlots(id, classes, dataset_dir, epoch, train_losses, train_accs,
test_losses, test_accs,
predict_list, true_list)
print('-' * 100)
print(f"TRAINING AND TESTING FINISHED")
print('-' * 100)
# train
print(f"TRAIN LOSS HISTORY: {train_losses}")
print('-' * 100)
print(f"TRAIN ACCURACY HISTORY: {train_accs}")
print('-' * 100)
# test
print('-' * 100)
print(f"TEST LOSS HISTORY: {test_losses}")
print('-' * 100)
print(f"TEST ACCURACY HISTORY: {test_accs}")
print('-' * 100)
print('-' * 100)
print('-' * 100)
print('-' * 100)
print("FINISH")
print('-' * 100)
print('-' * 100)
# 建立一个DataFrame对象，每一行包含每辆车的坐标和姿态信息。
train_images_dir = PATH + '/train_images/{}.jpg'
test_images_dir = PATH + '/test_images/{}.jpg'
df_train, df_dev = train_test_split(train, test_size=0.1,
random_state=42) # 划分训练集和验证集
df_test = test
df_dev_1 = df_dev.copy()
df_dev_pred_1 = pd.DataFrame()
df_dev_pred_2 = pd.DataFrame()
df_dev_pred_1['ImageId'] = df_dev_1['ImageId']
df_dev_pred_1['PredictionString'] = [''] * len(df_dev_1)
df_dev_pred_2['ImageId'] = df_dev_1['ImageId']
df_dev_pred_2['PredictionString'] = [''] * len(df_dev_1)
# Create dataset objects 创建数据集对象
train_dataset = CarDataset(df_train, train_images_dir, training=True)
dev_dataset = CarDataset(df_dev, train_images_dir, training=False)
test_dataset = CarDataset(df_test, test_images_dir, training=False)
# 数据集的设置
BATCH_SIZE = 4
train_loader = DataLoader(dataset=train_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=4)
dev_loader = DataLoader(dataset=dev_dataset,
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=0)
test_loader = DataLoader(dataset=test_dataset,
batch_size=BATCH_SIZE,
def train(args):
# basic arguments.
ngpu = args.ngpu
margin = args.margin
manual_seed = args.manual_seed
torch.manual_seed(manual_seed)
mean_value = args.mean_value
std_value = args.std_value
print("margin = {:5.2f}".format(margin))
print("manual_seed = {:5.2f}".format(manual_seed))
print("mean_value = {:5.2f}".format(mean_value))
print("std_value = {:5.2f}".format(std_value))
num_epochs = args.num_epochs
train_batch_size = args.train_batch_size
test_batch_size = args.test_batch_size
gamma = args.gamma # for learning rate decay
learning_rate = args.learning_rate
learning_rate2 = args.learning_rate2
loss_type = args.loss_type
dataset_name = args.dataset_name
pair_type = args.pair_type
mode = args.mode
weight_file = args.weight_file
print("pair_type = {}".format(pair_type))
print("loss_type = {}".format(loss_type))
print("mode = {}".format(mode))
print("weight_file = {}".format(weight_file))
root_dir = args.root_dir
image_txt = args.image_txt
train_test_split_txt = args.train_test_split_txt
label_txt = args.label_txt
ckpt_dir = args.ckpt_dir
eval_step = args.eval_step
display_step = args.display_step
embedding_size = args.embedding_size
pretrained = args.pretrained
aux_logits = args.aux_logits
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
kargs = {'ngpu': ngpu, 'pretrained': pretrained, 'aux_logits':aux_logits, 'embedding_size': embedding_size}
# create directory
model_dir = os.path.join(ckpt_dir, dataset_name, loss_type, str(int(embedding_size)))
print("model_dir = {}".format(model_dir))
if not os.path.isdir(model_dir):
os.makedirs(model_dir)
# network and loss
siamese_network = SiameseNetwork(**kargs)
first_group, second_group = siamese_network.separate_parameter_group()
param_lr_dict = [
{'params': first_group, 'lr': learning_rate2},
{'params': second_group, 'lr': learning_rate}
]
gpu_number = torch.cuda.device_count()
if device.type == 'cuda' and gpu_number > 1:
siamese_network = nn.DataParallel(siamese_network, list(range(torch.cuda.device_count())))
siamese_network.to(device)
# contrastive_loss = ContrastiveLoss(margin=margin)
# params = siamese_network.parameters()
print("args.optimizer = {:10s}".format(args.optimizer))
print("learning_rate = {:5.5f}".format(learning_rate))
print("learning_rate2 = {:5.5f}".format(learning_rate2))
optimizer = configure_optimizer(param_lr_dict, optimizer=args.optimizer)
# using different lr
# scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma, last_epoch=-1)
transform = transforms.Compose([transforms.Resize((299, 299)),
transforms.CenterCrop(299),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]
)
if dataset_name == 'cub200':
"""
print("dataset_name = {:10s}".format(dataset_name))
print(root_dir)
print(image_txt)
print(train_test_split_txt)
print(label_txt)
"""
dataset_train = CubDataset(root_dir, image_txt, train_test_split_txt, label_txt, transform=transform, is_train=True, offset=1)
dataset_eval = CubDataset(root_dir, image_txt, train_test_split_txt, label_txt, transform=transform, is_train=False, offset=1)
elif dataset_name == 'online_product':
"""
print("dataset_name = {:10s}".format(dataset_name))
"""
dataset_train = OnlineProductDataset(root_dir, train_txt=image_txt, test_txt=train_test_split_txt, transform=transform, is_train=True, offset=1)
dataset_eval = OnlineProductDataset(root_dir, train_txt=image_txt, test_txt=train_test_split_txt, transform=transform, is_train=False, offset=1)
elif dataset_name == "car196":
print("dataset_name = {}".format(dataset_name))
dataset_train = CarDataset(root_dir, image_info_mat=image_txt, transform=transform, is_train=True, offset=1)
dataset_eval = CarDataset(root_dir, image_info_mat=image_txt, transform=transform, is_train=False, offset=1)
dataloader = DataLoader(dataset=dataset_train, batch_size=train_batch_size, shuffle=False, num_workers=4)
dataloader_eval = DataLoader(dataset=dataset_eval, batch_size=test_batch_size, shuffle=False, num_workers=4)
log_for_loss = []
if mode == 'evaluation':
print("Do one time evluation and exit")
print("Load pretrained model")
siamese_network.module.load_state_dict(torch.load(weight_file))
print("Finish loading")
print("Calculting features")
feature_set, label_set, path_set = get_feature_and_label(siamese_network, dataloader_eval, device)
rec_pre = evaluation(feature_set, label_set)
# np.save("car196_rec_pre_ftl.npy", rec_pre)
# for visualization
sum_dict = {'feature': feature_set, 'label': label_set, 'path': path_set}
np.save('car196_fea_label_path.npy', sum_dict)
sys.exit()
print("Finish eval")
for epoch in range(num_epochs):
if epoch == 0:
feature_set, label_set, _ = get_feature_and_label(siamese_network, dataloader_eval, device)
# distance_type: Euclidean or cosine
rec_pre = evaluation(feature_set, label_set, distance_type='cosine')
siamese_network.train()
for i, data in enumerate(dataloader, 0):
# img_1, img_2, sim_label = data['img_1'].to(device), data['img_2'].to(device), data['sim_label'].type(torch.FloatTensor).to(device)
img_1, img_2, label_1, label_2 = data['img_1'].to(device), data['img_2'].to(device), data['label_1'].to(device), data['label_2'].to(device)
optimizer.zero_grad()
output_1, output_2 = siamese_network(img_1, img_2)
pair_dist, pair_sim_label = calculate_distance_and_similariy_label(output_1, output_2, label_1, label_2, sqrt=True, pair_type=pair_type)
if loss_type == "contrastive_loss":
loss, positive_loss, negative_loss = contrastive_loss(pair_dist, pair_sim_label, margin)
elif loss_type == "focal_contrastive_loss":
loss, positive_loss, negative_loss = focal_contrastive_loss(pair_dist, pair_sim_label, margin, mean_value, std_value)
elif loss_type == "triplet_loss":
loss, positive_loss, negative_loss = triplet_loss(pair_dist, pair_sim_label, margin)
elif loss_type == "focal_triplet_loss":
loss, positive_loss, negative_loss = focal_triplet_loss(pair_dist, pair_sim_label, margin, mean_value, std_value)
elif loss_type == "angular_loss":
center_output = (output_1 + output_2)/2.
pair_dist_2, _ = calculate_distance_and_similariy_label(center_output, output_2, label_1, label_2, sqrt=True, pair_type=pair_type)
# angle margin is 45^o
loss, positive_loss, negative_loss = angular_loss(pair_dist, pair_dist_2, pair_sim_label, 45)
else:
print("Unknown loss function")
sys.exit()
# try my own customized loss function
# loss = contrastive_loss(output_1, output_2, pair_sim_label)
loss.backward()
optimizer.step()
log_for_loss.append(loss.detach().item())
if i % display_step == 0 and i > 0:
print("{}, Epoch [{:3d}/{:3d}], Iter [{:3d}/{:3d}], Loss: {:6.5f}, Positive loss: {:6.5f}, Negative loss: {:6.5f}".format(
datetime.datetime.now(), epoch, num_epochs, i, len(dataloader), loss.item(), positive_loss.item(), negative_loss.item()))
if epoch % eval_step == 0:
print("Start evalution")
# np.save(loss_type +'.npy', log_for_loss)
feature_set, label_set, _ = get_feature_and_label(siamese_network, dataloader_eval, device)
# distance_type: Euclidean or cosine
rec_pre = evaluation(feature_set, label_set, distance_type='cosine')
torch.save(siamese_network.module.state_dict(), os.path.join(model_dir, 'model_' + str(epoch) +'_.pth'))
(can be specific for each pretrained neural network)
Return:
transform: albumentations.Compose
"""
_transform = [
albu.Lambda(image=preprocessing_fn),
albu.Lambda(image=to_tensor, mask=to_tensor),
]
return albu.Compose(_transform)
preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
test_dataset = CarDataset(img_dir_path=img_dir_path,
ano_dir_path=ano_dir_path,
aug=get_test_augmentation(),
preprocessing=get_preprocessing(preprocessing_fn),
predict=True)
print('num_test_images: {}'.format(test_dataset.__len__()))
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False)
# 可視化
def visualize(png_path, **images):
"""PLot images in one row."""
n = len(images)
plt.figure(figsize=(16, 5))
for i, (name, image) in enumerate(images.items()):