def create_dataloader(args):
train_transform = transforms.Compose([
# Resize(128),
transforms.Scale(256),
transforms.RandomCrop((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
test_transform = transforms.Compose([
# Resize(128),
transforms.Scale(256),
transforms.CenterCrop((224, 224)),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
train_path = os.path.join(args.data_path, 'train')
test_path = os.path.join(args.data_path, 'test')
train_dataset = dataset.ImageFolder(root=train_path, ground_truth=True, transform=train_transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4)
test_dataset = dataset.ImageFolder(root=test_path, ground_truth=False, transform=test_transform)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4)
return train_loader, test_loader
def get_data_file_list(root,
file_list,
batch_size,
train,
num_workers=0,
new_size=236,
crop_size=224):
if train:
transform = transforms.Compose([
transforms.Scale(new_size),
transforms.RandomSizedCrop(crop_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
else:
transform = transforms.Compose([
transforms.Resize((crop_size, crop_size)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
dataset = ImageLabelFilelist(root, file_list, transform=transform)
loader = DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=train,
drop_last=True,
num_workers=num_workers)
return loader
def main(args):
cfg_from_file(args.cfg)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
outputer = Outputer(args.output_dir, cfg.IMAGETEXT.PRINT_EVERY,
cfg.IMAGETEXT.SAVE_EVERY)
outputer.log(cfg)
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
val_transform = transforms.Compose([
transforms.Scale(int(imsize * 76 / 64)),
transforms.CenterCrop(imsize),
])
if args.dataset == 'bird':
val_dataset = ImageTextDataset(args.data_dir,
'test',
transform=val_transform,
sample_type='test')
elif args.dataset == 'coco':
val_dataset = CaptionDataset(args.data_dir,
'val',
transform=val_transform,
sample_type='test')
else:
raise NotImplementedError
val_dataloader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=int(cfg.WORKERS))
if args.raw_checkpoint != '':
outputer.log("load raw checkpoint from {}".format(args.raw_checkpoint))
encoder, decoder = load_raw_checkpoint(args.raw_checkpoint)
else:
# define the model
encoder = Encoder()
encoder.fine_tune(False)
decoder = DecoderWithAttention(
attention_dim=cfg.IMAGETEXT.ATTENTION_DIM,
embed_dim=cfg.IMAGETEXT.EMBED_DIM,
decoder_dim=cfg.IMAGETEXT.DECODER_DIM,
vocab_size=val_dataset.n_words)
assert args.model_path != ""
outputer.log("load model dict from {}".format(args.model_path))
encoder, decoder = load_checkpoint(encoder, decoder, args.model_path)
encoder.fine_tune(False)
encoder, decoder = encoder.cuda(), decoder.cuda()
outputer.log("eval the model...")
pred_file = os.path.join(outputer.log_folder, 'pred_result.csv')
eval_rtn, outputer = evaluate(val_dataloader, encoder, decoder, outputer,
args.beam_size, pred_file)
outputer.log("eval result: {}".format(dict2str(eval_rtn)))
def __init__(self, N, train=True, seed=42, device='cpu', noise_std=0.1):
np.random.seed(seed)
self.IMG_WIDTH, self.IMG_HEIGHT = 28, 28
self.mnist = datasets.MNIST('../data',
train=train,
download=True,
transform=transforms.Compose([
transforms.Scale(self.IMG_WIDTH),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
]))
self.n = len(self.mnist)
self.N = N
self.device = device
self.noise_std = noise_std
def get_transform(args):
if args.dataset == 'celeba':
crop_size = 108
re_size = 64
offset_height = (218 - crop_size) // 2
offset_width = (178 - crop_size) // 2
crop = lambda x: x[:, offset_height:offset_height + crop_size,
offset_width:offset_width + crop_size]
preprocess = transforms.Compose(
[transforms.ToTensor(),
transforms.Lambda(crop),
transforms.ToPILImage(),
transforms.Scale(size=(re_size, re_size), interpolation=Image.BICUBIC),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)])
return preprocess
model_ft.fc = torch.nn.Linear(num_ftrs, 2)
learning_rate = 0.001
criterion = torch.nn.CrossEntropyLoss()
optimizer_ft = torch.optim.SGD(model_ft.parameters(),
lr=0.001,
momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft,
step_size=7,
gamma=0.1)
if is_cuda:
model_ft = model_ft.cuda()
print("model_ft", model_ft)
simple_transform = transforms.Compose([
transforms.Scale((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# p = "C:/Users/a/Downloads/train/train/dog"
train = ImageFolder("C:/Users/a/Downloads/train/train/", simple_transform)
valid = ImageFolder("C:/Users/a/Downloads/train/valid/", simple_transform)
train_data_gen = torch.utils.data.DataLoader(train,
batch_size=64,
num_workers=3)
valid_data_gen = torch.utils.data.DataLoader(valid,
batch_size=64,
num_workers=3)
def Scale(self, **args):
return self._add(transforms.Scale(**args))
snapshot_path = args.snapshot
# ResNet50 structure
model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3],
66)
print 'Loading snapshot.'
# Load snapshot
save_state_dict = torch.load(snapshot_path)
model.load_state_dict(save_state_dict)
print 'Loading data.'
transformations = transforms.Compose([
transforms.Scale(140),
transforms.CenterCrop(128),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
if args.dataset == 'Driver_Test':
print args.data_dir, args.filename_list
pose_dataset = datasets.Driver_Test(args.data_dir, args.filename_list,
transformations)
elif args.dataset == 'By_TXT':
pose_dataset = datasets.Driver_Train_TXT(args.filename_list,
transformations)
test_loader = torch.utils.data.DataLoader(dataset=pose_dataset,
def train():
# Define training hyperparameters
epochs = 5
learning_rate = 0.02
log_interval = 10
training_batch_size = 128
testing_batch_size = 1000
crop_size = 28
channels = 3
image_size = 64
#Adam optimizer parameter
beta1 = 0.5
# number of gpus
ngpu = 1
num_epochs = 5
# Decide which device we want to run on
#device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
# Create the generator
generator = Generator(128)
discriminator = Discriminator(128)
#if (device.type == 'cuda') and (ngpu > 1):
# generator = nn.DataParallel(generator, list(range(ngpu)))
# discriminator = nn.DataParallel(discriminator, list(range(ngpu)))
# Apply the weights_init function to randomly initialize all weights
# to mean=0, stdev=0.2.
generator.apply(weights_init)
discriminator.apply(weights_init)
transform = transforms.Compose([
transforms.Scale(image_size),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])
# Load Training Data...
training_data_path = 'cars_train'
#NOTE: Comment either out for training on a specific dataset
#dataset = torchvision.datasets.ImageFolder('cars_train',transform=transform)
dataset = torchvision.datasets.MNIST('./',
train=True,
transform=transform,
target_transform=None,
download=True)
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=training_batch_size,
shuffle=True,
num_workers=4)
# Plot some training images
real_batch = next(iter(data_loader))
plt.figure(figsize=(8, 8))
plt.axis("off")
plt.title("Training Images")
plt.imshow(
np.transpose(
vutils.make_grid(real_batch[0][:64], padding=2,
normalize=True).cpu(), (1, 2, 0)))
#plt.show()
size = len(dataset)
# Initialize BCELoss function
criterion = nn.BCEWithLogitsLoss()
# Create batch of latent vectors
fixed_noise = torch.randn(64, 100, 1, 1)
real_label = 1
fake_label = 0
# Setup Adam optimizers for both G and D
D_optimizer = optim.Adam(discriminator.parameters(),
lr=learning_rate,
betas=(beta1, 0.999))
G_optimizer = optim.Adam(generator.parameters(),
lr=learning_rate,
betas=(beta1, 0.999))
# Lists to keep track of progress
img_list = []
G_losses = []
D_losses = []
iters = 0
print("Starting Training Loop...")
for epoch in range(num_epochs):
for i, data in enumerate(data_loader, 0):
## Train with real data batch
discriminator.zero_grad()
real_cpu = data[0]
b_size = real_cpu.size(0)
l = list(data[1].size())[0]
output_real = torch.ones(l)
output_fake = torch.zeros(l)
output = discriminator(data[0]).view(-1)
d_real_image_error = criterion(output, output_real)
# Calculate gradients for D in backward pass
d_real_image_error.backward()
D_x = output.mean().item()
#noise vector
noise = torch.randn(b_size, 100, 1, 1)
# Generate fake image batch with G
fake = generator(noise)
#label.fill_(fake_label)
# Classify all fake batch with D
output = discriminator(fake.detach()).view(-1)
# Calculate D's loss on the all-fake batch
d_fake_image_error = criterion(output, output_fake)
# Calculate the gradients for this batch
d_fake_image_error.backward()
D_G_z1 = output.mean().item()
errD = d_real_image_error + d_fake_image_error
# Update D
D_optimizer.step()
#Update the Generator
generator.zero_grad()
#label.fill_(label)
output = discriminator(fake).view(-1)
generator_err = criterion(output, output_real)
# Calculate gradients for G
generator_err.backward()
D_G_z2 = output.mean().item()
G_optimizer.step()
# Save Losses for plotting later
G_losses.append(generator_err.item())
D_losses.append(errD.item())
iters += 1
print("iters: " + str(iters))
stamp = datetime.now().strftime("%m%d%Y_%H%M%S")
save_network(generator, discriminator, stamp, epoch)
def main(args):
cfg_from_file(args.cfg)
if args.gpu_id != -1:
cfg.GPU_ID = args.gpu_id
else:
cfg.CUDA = False
if args.data_dir != '':
cfg.DATA_DIR = args.data_dir
if args.textimage_nete_path != '':
cfg.TRAIN.NET_E = args.textimage_nete_path
pprint(cfg)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
outputer = Outputer(args.output_dir, cfg.IMAGETEXT.PRINT_EVERY,
cfg.IMAGETEXT.SAVE_EVERY)
outputer.log(cfg)
imsize = cfg.TREE.BASE_SIZE * (2**(cfg.TREE.BRANCH_NUM - 1))
val_transform = transforms.Compose([
transforms.Scale(int(imsize * 76 / 64)),
transforms.CenterCrop(imsize),
])
if args.dataset == 'bird':
imagetext_val_dataset = ImageTextDataset(args.data_dir,
'test',
transform=val_transform,
sample_type='test')
elif args.dataset == 'coco':
imagetext_val_dataset = CaptionDataset(args.data_dir,
'val',
transform=val_transform,
sample_type='test')
else:
raise NotImplementedError
imagetext_val_dataloader = torch.utils.data.DataLoader(
imagetext_val_dataset,
batch_size=1,
shuffle=False,
num_workers=int(cfg.WORKERS))
# evaluation imagetext model
if args.imagetext_model_path != '':
outputer.log("load raw checkpoint from: {}".format(
args.imagetext_model_path))
if os.path.splitext(
args.imagetext_model_path)[-1] == '.tar': # raw checkpoint
imagetext_encoder, imagetext_decoder = load_raw_checkpoint(
args.imagetext_model_path)
else: #define the model
imagetext_encoder = Encoder()
imagetext_encoder.fine_tune(False)
imagetext_decoder = DecoderWithAttention(
attention_dim=cfg.IMAGETEXT.ATTENTION_DIM,
embed_dim=cfg.IMAGETEXT.EMBED_DIM,
decoder_dim=cfg.IMAGETEXT.DECODER_DIM,
vocab_size=imagetext_val_dataset.n_words)
imagetext_encoder, imagetext_decoder = load_checkpoint(
imagetext_encoder, imagetext_decoder,
args.imagetext_model_path)
else:
print("the model path for imagetext model is NULL")
raise ValueError
imagetext_encoder, imagetext_decoder = imagetext_encoder.cuda(
), imagetext_decoder.cuda()
text_file = os.path.join(outputer.log_folder, "pred_result.csv")
imagetext_eval_rtn, outputer = evaluate(imagetext_val_dataloader,
imagetext_encoder,
imagetext_decoder, outputer,
args.beam_size, text_file)
# write the pred to local file
text_list = []
filename_list = []
with open(text_file, 'r') as fp:
for _l in fp.readlines():
data_list = _l.strip().split(',')
text_list.append(data_list[-2])
filename_list.append(data_list[-1])
#--------- generate images from the text -------#
if args.dataset == 'bird':
textimage_val_dataset = TextDataset(args.data_dir,
'test',
base_size=cfg.TREE.BASE_SIZE,
transform=val_transform)
elif args.dataset == 'coco':
textimage_val_dataset = CaptionDatasetMultisize(
args.data_dir,
'val',
base_size=cfg.TREE.BASE_SIZE,
transform=val_transform)
else:
raise ValueError
textimage_val_dataloader = torch.utils.data.DataLoader(
textimage_val_dataset, batch_size=1)
image_output_dir = os.path.join(args.output_dir, "images")
algo = trainer(image_output_dir, textimage_val_dataloader,
textimage_val_dataset.n_words,
textimage_val_dataset.ixtoword)
gen_example(textimage_val_dataset.wordtoix, algo, text_list, filename_list)
def main():
# img_folders = ["../coco_img/bg2_0/", "../coco_img/bg2_127/", "../coco_img/bg2_255/",
# "../coco_img/obj2_0/", "../coco_img/obj2_127/", "../coco_img/obj2_255/"]
img_folders = [
"../coco_img/merged_bg2_0/", "../coco_img/merged_bg2_127/",
"../coco_img/merged_bg2_255/", "../coco_img/merged_obj2_0/",
"../coco_img/merged_obj2_127/", "../coco_img/merged_obj2_255/"
]
img_folders = ["../coco_img/org/"]
model_name = "MLCCOCO"
model = MultilabelObject(None, 80).cuda()
log_dir = "./"
checkpoint = torch.load(os.path.join(log_dir, 'model_best.pth.tar'),
encoding='bytes')
new_checkpoint = OrderedDict()
for k in checkpoint[b'state_dict']:
new_checkpoint[k.decode('utf-8')] = checkpoint[b'state_dict'][k]
model.load_state_dict(new_checkpoint)
model.eval()
with open("classes_list.pickle", "rb") as f:
classes_list = pickle.load(f)
for img_folder in img_folders:
crop_size = 224
image_size = 256
batch_size = 64
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
val_transform = transforms.Compose([
transforms.Scale(image_size),
transforms.CenterCrop(crop_size),
transforms.ToTensor(), normalize
])
cocomlc = COCO_MLC(img_folder, val_transform)
test_loader = torch.utils.data.DataLoader(cocomlc,
batch_size=batch_size,
shuffle=False,
num_workers=8,
pin_memory=True)
t = tqdm(test_loader, desc='testing {}'.format(img_folder))
result = {}
for batch_idx, (imgs, paths) in enumerate(t):
images = imgs.cuda()
# print(images.shape)
probs, labels, labels_probs = infer_batch(model,
classes_list,
inputs=images,
threshold=0.5)
for i in range(0, len(paths)):
path = paths[i]
result[path] = {
"prob": probs[i],
"labels": labels[i],
"labels_probs": labels_probs[i]
}
pickle_file_name = "{}_{}.pickle".format(
model_name, os.path.basename(os.path.normpath(img_folder)))
pickle_path = os.path.join(".", "result_pickle", pickle_file_name)
with open(pickle_path, 'wb') as handle:
pickle.dump(result, handle)
print("Done, Saved to {}".format(pickle_path))
def train(nb_epoch,
batch_size,
store_name,
image_path,
soft_labels=False,
gaussian_std=1.5,
resume=False,
start_epoch=0,
model_path=None):
batch_size = int(batch_size)
nb_epoch = int(nb_epoch)
start_epoch = int(start_epoch)
exp_dir = store_name
try:
os.stat(exp_dir)
except:
os.makedirs(exp_dir)
use_cuda = torch.cuda.is_available()
# GPU
# Data
print('==> Preparing data..')
transform_train = transforms.Compose([
transforms.Scale((550, 550)),
transforms.RandomCrop(448, padding=8),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
transform_test = transforms.Compose([
transforms.Scale((550, 550)),
transforms.CenterCrop(448),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
trainset = YearBuiltFolder(image_path,
calc_perf=True,
soft_labels=soft_labels,
gaussian_std=gaussian_std,
transforms=transform_train)
train_weights = np.array(trainset.train_weights)
train_sampler = torch.utils.data.WeightedRandomSampler(train_weights,
len(train_weights),
replacement=True)
if use_cuda:
device = 'cuda'
else:
device = 'cpu'
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=(train_sampler == None),
num_workers=4,
sampler=train_sampler)
# Testing is on the same data but without the transforms. This whole class
# only givey training performance and the user has to use the detect script
# to check the performance.
train_set_without_transforms = YearBuiltFolder(image_path,
calc_perf=True,
soft_labels=soft_labels,
gaussian_std=gaussian_std,
transforms=transform_test)
train_test_loader = torch.utils.data.DataLoader(
train_set_without_transforms,
batch_size=3,
shuffle=True,
num_workers=4)
# Model
if resume:
net = torch.load(model_path)
else:
net = load_model(model_name='resnet50_pmg',
pretrain=True,
require_grad=True,
num_classes=len(trainset.classes))
netp = torch.nn.DataParallel(net)
net.to(device)
if soft_labels:
loss = nn.KLDivLoss(reduction='batchmean')
else:
loss = nn.CrossEntropyLoss()
sm = nn.LogSoftmax(dim=1)
optimizer = optim.SGD([{
'params': net.classifier_concat.parameters(),
'lr': 0.002
}, {
'params': net.conv_block1.parameters(),
'lr': 0.002
}, {
'params': net.classifier1.parameters(),
'lr': 0.002
}, {
'params': net.conv_block2.parameters(),
'lr': 0.002
}, {
'params': net.classifier2.parameters(),
'lr': 0.002
}, {
'params': net.conv_block3.parameters(),
'lr': 0.002
}, {
'params': net.classifier3.parameters(),
'lr': 0.002
}, {
'params': net.features.parameters(),
'lr': 0.0002
}],
momentum=0.9,
weight_decay=5e-4)
max_train_acc = 0
lr = [0.002, 0.002, 0.002, 0.002, 0.002, 0.002, 0.002, 0.0002]
for epoch in range(start_epoch, nb_epoch):
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
train_loss1 = 0
train_loss2 = 0
train_loss3 = 0
train_loss4 = 0
correct = 0
total = 0
idx = 0
for batch_idx, (inputs, targets, indexes) in enumerate(trainloader):
idx = batch_idx
if inputs.shape[0] < batch_size:
continue
if use_cuda:
inputs, targets = inputs.to(device), targets.to(device)
inputs, targets = Variable(inputs), Variable(targets)
# update learning rate
for nlr in range(len(optimizer.param_groups)):
optimizer.param_groups[nlr]['lr'] = cosine_anneal_schedule(
epoch, nb_epoch, lr[nlr])
# Step 1
optimizer.zero_grad()
inputs1 = jigsaw_generator(inputs, 8)
output_1, _, _, _ = netp(inputs1)
loss1 = loss(sm(output_1), targets) * 1
loss1.backward()
optimizer.step()
# Step 2
optimizer.zero_grad()
inputs2 = jigsaw_generator(inputs, 4)
_, output_2, _, _ = netp(inputs2)
loss2 = loss(sm(output_2), targets) * 1
loss2.backward()
optimizer.step()
# Step 3
optimizer.zero_grad()
inputs3 = jigsaw_generator(inputs, 2)
_, _, output_3, _ = netp(inputs3)
loss3 = loss(sm(output_3), targets) * 1
loss3.backward()
optimizer.step()
# Step 4
optimizer.zero_grad()
_, _, _, output_concat = netp(inputs)
concat_loss = loss(sm(output_concat), targets) * 2
concat_loss.backward()
optimizer.step()
# training log
_, predicted = torch.max(output_concat.data, 1)
if soft_labels:
_, targets = torch.max(targets.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum()
train_loss += (loss1.item() + loss2.item() + loss3.item() +
concat_loss.item())
train_loss1 += loss1.item()
train_loss2 += loss2.item()
train_loss3 += loss3.item()
train_loss4 += concat_loss.item()
if batch_idx % 50 == 0:
print(
'Step: %d | Loss1: %.3f | Loss2: %.5f | Loss3: %.5f | Loss_concat: %.5f | Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (batch_idx, train_loss1 / (batch_idx + 1), train_loss2 /
(batch_idx + 1), train_loss3 /
(batch_idx + 1), train_loss4 /
(batch_idx + 1), train_loss / (batch_idx + 1),
100. * float(correct) / total, correct, total))
train_acc = 100. * float(correct) / total
train_loss = train_loss / (idx + 1)
with open(os.path.join(exp_dir, 'results_train.txt'), 'a') as file:
file.write(
'Iteration %d | train_acc = %.5f | train_loss = %.5f | Loss1: %.3f | Loss2: %.5f | Loss3: %.5f | Loss_concat: %.5f |\n'
% (epoch, train_acc, train_loss, train_loss1 /
(idx + 1), train_loss2 / (idx + 1), train_loss3 /
(idx + 1), train_loss4 / (idx + 1)))
# if epoch < 5 or epoch >= 80:
print('Start testing')
if soft_labels:
train_acc, train_acc_com, train_loss, y_gt, y_pred = test_softlabels(
net, loss, train_test_loader)
else:
train_acc, train_acc_com, train_loss, y_gt, y_pred = test(
net, loss, train_test_loader)
if train_acc_com > max_train_acc:
max_train_acc = train_acc_com
net.cpu()
torch.save(net, os.path.join(store_name, 'model_best.pth'))
net.to(device)
with open(os.path.join(exp_dir, 'results_test.txt'), 'a') as file:
file.write(
'Iteration %d, test_acc = %.5f, test_acc_combined = %.5f, test_loss = %.6f\n'
% (epoch, train_acc, train_acc_com, train_loss))
net.cpu()
torch.save(
net, os.path.join(store_name, 'model_epoch_{}.pth'.format(epoch)))
net.to(device)
import cv2
import torch as t
import numpy as np
from torch.tensor import Tensor
from torchvision.transforms import transforms
from torchvision.transforms import ToPILImage
to_pil_image = ToPILImage()
transform = transforms.Compose([
transforms.ToTensor(), # 转为Tensor
transforms.Scale((32, 32)), # resize
# transforms.CenterCrop((32, 32)),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), # 归一化
])
# todo: 训练模型
# 1. 构造训练集
# 2. 读取训练集
# 3. 算法（网络定义）
# 4. 保存模型文件，输出验证集结果
def read_image(image_path):
"""读出图片数据"""
img = transform(cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_RGB2BGR))
img = np.reshape(img, (3, 32, 32))
# debug: 输出图片
# img = to_pil_image(img)
# img.show()
return Tensor(img).unsqueeze(0)
#def nms(dets, thresh):
# "Dispatch to either CPU or GPU NMS implementations.\
# Accept dets as tensor"""
# return pth_nms(dets, thresh)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
IMG_SIZE = (480, 640)
HEATMAP_SIZE = (29, 39)
TEMPLATE_SIZE = 124
PREPROCESS = [
transforms.Compose(
[transforms.Scale(IMG_SIZE[0]),
transforms.ToTensor(), normalize]),
transforms.Compose(
[transforms.Scale(TEMPLATE_SIZE),
transforms.ToTensor(), normalize]),
transforms.Compose(
[transforms.Scale(TEMPLATE_SIZE),
transforms.ToTensor()]),
transforms.Compose(
[transforms.Scale(TEMPLATE_SIZE),
transforms.ToTensor(), normalize]),
transforms.Compose(
[transforms.Scale(TEMPLATE_SIZE),
transforms.ToTensor()])
]
