def main():
print("loading data...")
ds = dataset.Dataset(classes=classes)
train_X, train_y = ds.load_data('train')
train_X = ds.preprocess_inputs(train_X)
train_Y = ds.reshape_labels(train_y)
print("input data shape...", train_X.shape)
print("input label shape...", train_Y.shape)
test_X, test_y = ds.load_data('test')
test_X = ds.preprocess_inputs(test_X)
test_Y = ds.reshape_labels(test_y)
print("creating model...")
model = SegNet(input_shape=input_shape, classes=classes)
model.compile(loss="categorical_crossentropy",
optimizer='adadelta',
metrics=["accuracy"])
model.fit(train_X,
train_Y,
batch_size=batch_size,
epochs=epochs,
verbose=1,
class_weight=class_weighting,
validation_data=(test_X, test_Y),
shuffle=True)
model.save('seg.h5')
def main():
# Parse arguments.
parser = argparse.ArgumentParser()
kwargs = {
'type': str,
'help': 'The model file path.',
'required': True
}
parser.add_argument('-m', '--model', **kwargs)
kwargs = {
'type': int,
'default': 10,
'help': 'The number of samples to evaluate. default: 10'
}
parser.add_argument('-n', '--num', **kwargs)
kwargs = {
'type': str,
'default': 'val',
'help': 'Type of dataset to use. "val" or "test". default: "val"'
}
parser.add_argument('-t', '--type', **kwargs)
args = parser.parse_args()
# Prepare training data.
dataset = np.load('./temp/dataset.npz')
val_x = dataset[f'{args.type}_x']
val_y = dataset[f'{args.type}_y']
if args.num < 0 or len(val_x) < args.num:
args.num = len(val_x)
# Prepare tensorflow.
config = tf.ConfigProto(device_count={'GPU': 0})
session = tf.Session(config=config)
keras.backend.tensorflow_backend.set_session(session)
# Prepare model.
model = SegNet(shape=(360, 480, 3))
model.load_weights(args.model)
# Output results.
head, tail = os.path.split(args.model)
filename, ext = os.path.splitext(tail)
os.makedirs(f'{head}/{filename}/', exist_ok=True)
for i, x, y, t in zip(range(args.num), val_x, model.predict(val_x[:args.num]), val_y):
cv2.imwrite(f'{head}/{filename}/{args.type}-{i}-input.png', x * 255)
cv2.imwrite(f'{head}/{filename}/{args.type}-{i}-prediction.png', y * 255)
cv2.imwrite(f'{head}/{filename}/{args.type}-{i}-teacher.png', t * 255)
def main():
input_shape = (360, 480, 3)
classes = 12
epochs = 100
batch_size = 1
log_path = './logs/'
class_weighting = [
0.2595, 0.1826, 4.5640, 0.1417, 0.5051, 0.3826, 9.6446, 1.8418, 6.6823,
6.2478, 3.0, 7.3614
]
# set gpu usage
config = tf.ConfigProto(gpu_options=tf.GPUOptions(
allow_growth=True, per_process_gpu_memory_fraction=0.8))
session = tf.Session(config=config)
set_session(session)
print("loading data...")
ds = dataset.Dataset(classes=classes,
train_file="CamVid/train.txt",
test_file="CamVid/test.txt")
# need to implement, y shape is (None, 360, 480, classes)
train_x, train_y = ds.load_data(root_path="CamVid", mode='train')
train_x = ds.preprocess_inputs(train_x)
train_y = ds.reshape_labels(train_y)
print("input data shape...", train_x.shape)
print("input label shape...", train_y.shape)
# need to implement, y shape is (None, 360, 480, classes)
test_x, test_y = ds.load_data(root_path="CamVid", mode='test')
test_x = ds.preprocess_inputs(test_x)
test_y = ds.reshape_labels(test_y)
tb_cb = TensorBoard(log_dir=log_path,
histogram_freq=1,
write_graph=True,
write_images=True)
print("creating model...")
model = SegNet(input_shape=input_shape, classes=classes)
model.compile(loss="categorical_crossentropy",
optimizer='adadelta',
metrics=["accuracy"])
model.fit(train_x,
train_y,
batch_size=batch_size,
epochs=epochs,
verbose=1,
class_weight=class_weighting,
validation_data=(test_x, test_y),
shuffle=True,
callbacks=[tb_cb])
model.save('seg.h5')
def main():
model_dir = './checkpoints/seg2/segnet_gen1/model_at_epoch_013.dat'
save_dir = './test/0610/segnet_gen1/test'
test_txt_path = './data/seg/valid.txt'
# model = unet(in_channel=1, n_classes=1)
model = SegNet(input_nbr = 1, label_nbr = 1)
model = load_model(model, model_dir)
model = model.cuda()
model.eval()
test_dataset = GuideWireDataset(test_txt_path)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS)
prefetcher = data_prefetcher(test_loader)
input, target, distance = prefetcher.next()
dice = []
IoU = []
precision = []
recall = []
i = -1
while input is not None:
i += 1
with torch.no_grad():
output = model(input)
dice.append(dice_coeff(output, target).item())
IoU.append(iou_coeff(output, target).item())
precision.append(Precision(output, target).item())
recall.append(Recall(output, target).item())
output = torch.sigmoid(output).squeeze().data.cpu().numpy()
output[output < 0.5] = 0
output[output >= 0.5] = 1
# output = torch.argmax(output, dim=1).squeeze().data.cpu().numpy()
# output = output.squeeze().data.cpu().numpy()
# output = np.argmax(output, axis=0)
cv2.imwrite(os.path.join(save_dir, str(i) + '_output.jpg'), output * 255)
print(str(i) + ' finish!')
input, target, distance = prefetcher.next()
print('dice: ', np.mean(dice), np.max(dice), np.min(dice), np.std(dice))
print('iou: ', np.mean(IoU), np.max(IoU), np.min(IoU), np.std(IoU))
print('precision: ', np.mean(precision), np.max(precision), np.min(precision), np.std(precision))
print('recall: ', np.mean(recall), np.max(recall), np.min(recall), np.std(recall))
def main():
print("loading data...")
ds = dataset.DataSet(classes=classes)
train_X, train_y = ds.load_data(
'train') # need to implement, y shape is (None, 360, 480, classes)
train_X = ds.preprocess_inputs(train_X)
train_Y = ds.reshape_labels(train_y)
print("input data shape...", train_X.shape)
print("input label shape...", train_Y.shape)
test_X, test_y = ds.load_data(
'test') # need to implement, y shape is (None, 360, 480, classes)
test_X = ds.preprocess_inputs(test_X)
test_Y = ds.reshape_labels(test_y)
tb_cb = keras.callbacks.TensorBoard(log_dir=log_filepath,
histogram_freq=1,
write_graph=True,
write_images=True)
fpath = 'weights.{epoch:02d}.hdf5'
mc_cb = keras.callbacks.ModelCheckpoint(fpath,
monitor='val_loss',
verbose=0,
save_best_only=False,
save_weights_only=False,
mode='auto',
period=3)
print("creating model...")
model = SegNet(input_shape=input_shape, classes=classes)
model.compile(loss="categorical_crossentropy",
optimizer='adadelta',
metrics=["accuracy"])
model.fit(train_X,
train_Y,
batch_size=batch_size,
epochs=epochs,
verbose=1,
class_weight=class_weighting,
validation_data=(test_X[0:5], test_Y[0:5]),
shuffle=True,
callbacks=[mc_cb, tb_cb])
model.save('seg.h5')
def test(img_path, data_size='single'):
device = torch.device('cuda') if torch.cuda.is_available else torch.device(
'cpu')
# Image input
im = Image.open(img_path)
im = np.array(im, dtype=np.float32) / 255
image = np.transpose(im, (2, 0, 1))
data = torch.from_numpy(image).unsqueeze(0)
data = Variable(data).to(device)
model = SegNet(opt, data.shape[1])
if opt.model_path:
model.load_state_dict(torch.load(opt.model_path))
model = model.to(device)
model.train()
feats, output = model(data)
output = output[0].permute(1, 2, 0).contiguous().view(-1, opt.nClass)
feats = feats[0].permute(1, 2, 0).contiguous().view(-1, opt.nChannel)
_, pred_clusters = torch.max(output, 1)
pred_clusters = pred_clusters.data.cpu().numpy()
# Post processing
labels = np.unique(pred_clusters)
counts = {}
for i in pred_clusters:
counts[i] = counts.get(i, 0) + 1
sorts = sorted(counts.items(), key=lambda x: x[1])
cache = {}
cache[sorts[-1][0]] = 0
n = 1
for (num, _) in sorts[:-1]:
cache[num] = n
n += 1
label_colors = [[10, 10, 10], [0, 0, 255], [0, 255, 0], [255, 0, 0],
[255, 255, 0], [0, 255, 255], [255, 0, 255]]
im_target_rgb = np.array([label_colors[cache[c]] for c in pred_clusters])
im_target_rgb = im_target_rgb.reshape(im.shape).astype(np.uint8)
# change path
path = ".".join(img_path.split('/')[1].split('.')[:2])
#path = img_path.split('/')[1].split('.')[0]
if data_size == 'single':
cv2.imwrite("outputs_single/{}_out.png".format(path), im_target_rgb)
elif data_size == 'all':
cv2.imwrite("outputs_all/{}_out.png".format(path), im_target_rgb)
def test_autoencoder(epoch_plus, text, index):
use_gpu = torch.cuda.is_available()
ngpu = torch.cuda.device_count()
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
model = SegNet(3)
if ngpu > 1:
model = nn.DataParallel(model)
if use_gpu:
model = model.to(device, non_blocking=True)
text = text.to(device, non_blocking=True)
if epoch_plus > 0:
model.load_state_dict(torch.load('./autoencoder_models_2/autoencoder_{}.pth'.format(epoch_plus)))
model.eval()
if use_gpu:
text.to(device, non_blocking=True)
predicted = model(text)
predicted[predicted > 1.0] = 1.0
save_path1 = './results/text'
save_path2 = './results/masked'
if not os.path.exists(save_path1):
os.mkdir(save_path1)
if not os.path.exists(save_path2):
os.mkdir(save_path2)
binary_predicted = predicted.clone()
binary_mask = predicted.clone()
binary_predicted[binary_predicted > 0.0] = 1.0
binary_mask[binary_mask > 0.1] = 1.0
masked = text + binary_mask
masked[masked > 1.0] = 1.0
trans = torchvision.transforms.ToPILImage()
predicted = predicted.squeeze().cpu()
masked = masked.squeeze().cpu()
image = trans(predicted)
image2 = trans(masked)
image.save(os.path.join(save_path1, 'text_{}.png'.format(index)))
image2.save(os.path.join(save_path2, 'masked_{}.png'.format(index)))
del text
del predicted
del masked
del binary_predicted
def train(data_size='all'):
device = torch.device('cuda') if torch.cuda.is_available else torch.device(
'cpu')
# Image input
model = SegNet(opt, 3)
model = model.to(device)
model.train()
criterion = torch.nn.CrossEntropyLoss()
criterion_d = DiscriminativeLoss()
optimizer = SGD(model.parameters(), lr=opt.lr, momentum=opt.momentum)
if data_size == 'all':
dataloader = get_dataloader(opt.paths, opt, device)
model = batch_step(opt, optimizer, model, dataloader, criterion,
criterion_d, device)
torch.save(model.state_dict(), 'model_all.pth')
else:
im = Image.open(opt.img_path)
im = np.array(im, dtype=np.float32) / 255
image = np.transpose(im, (2, 0, 1))
data = torch.from_numpy(image).unsqueeze(0)
data = Variable(data).to(device)
labels = segmentation.slic(im,
compactness=opt.compactness,
n_segments=opt.num_superpixels)
labels = labels.reshape(-1)
label_nums = np.unique(labels)
label_indices = [
np.where(labels == label_nums[i])[0]
for i in range(len(label_nums))
]
model = one_step(opt, optimizer, model, data, label_indices, criterion,
criterion_d, device)
torch.save(model.state_dict(), 'model_single.pth')
def main():
print("loading data...")
ds = dataset.Dataset(test_file='test_5.txt', classes=classes)
train_X, train_y = ds.load_data('train') # need to implement, y shape is (None, 360, 480, classes)
train_X = ds.preprocess_inputs(train_X)
train_Y = ds.reshape_labels(train_y)
print("input data shape...", train_X.shape)
print("input label shape...", train_Y.shape)
test_X, test_y = ds.load_data('test') # need to implement, y shape is (None, 360, 480, classes)
test_X = ds.preprocess_inputs(test_X)
test_Y = ds.reshape_labels(test_y)
tb_cb = keras.callbacks.TensorBoard(log_dir=log_filepath, histogram_freq=1, write_graph=True, write_images=True)
print("creating model...")
model = SegNet(input_shape=input_shape, classes=classes)
model.compile(loss="categorical_crossentropy", optimizer='adadelta', metrics=["accuracy"])
model.fit(train_X, train_Y, batch_size=batch_size, epochs=epochs,
verbose=1, class_weight=class_weighting , validation_data=(test_X, test_Y), shuffle=True
, callbacks=[tb_cb])
model.save('s.h5')
def predict_image(dir):
use_gpu = torch.cuda.is_available()
ngpu = torch.cuda.device_count()
device = torch.device("cuda:0" if (
torch.cuda.is_available() and ngpu > 0) else "cpu")
image_to_tensor = torchvision.transforms.ToTensor()
tensor_to_image = torchvision.transforms.ToPILImage()
save_path = Path(dir).parent
image = Image.open(dir).convert('RGB')
image = image_to_tensor(image)
c, w, h = image.shape
image = torch.reshape(image, (1, c, w, h))
model = SegNet(3)
if use_gpu:
model = model.to(device, non_blocking=True)
image = image.to(device, non_blocking=True)
model.load_state_dict(torch.load('./models/model.pth',
map_location=device))
model.eval()
predicted = model(image)
predicted[predicted > 1.0] = 1.0
binary_predicted = predicted.clone()
binary_mask = predicted.clone()
binary_predicted[binary_predicted > 0.0] = 1.0
binary_mask[binary_mask > 0.1] = 1.0
masked = image + binary_mask
masked[masked > 1.0] = 1.0
predicted = predicted.squeeze().cpu()
masked = masked.squeeze().cpu()
image = tensor_to_image(predicted)
image2 = tensor_to_image(masked)
image.save(os.path.join(save_path, 'tmp_text.png'))
image2.save(os.path.join(save_path, 'tmp_masked.png'))
def test(args):
cfg = load_cfg(args.cfg)
weight_path = args.wts
img_path = args.im_path
segnet = SegNet().float().cuda()
segnet.load_state_dict(torch.load(weight_path))
segnet.eval()
im = cv2.imread(img_path).transpose(2, 0, 1)
im = torch.tensor(im[np.newaxis, :], dtype=torch.float).cuda()
out = segnet(im)
out = out.detach().cpu().numpy().transpose(0, 2, 3, 1)
out = np.argmax(out, axis=3).astype(np.uint8)[0]
out = out[:, :, np.newaxis]
out = out * 20
cv2.imshow('f**k', out)
cv2.waitKey(0)
plt.ylabel("Loss")
plt.show()
plt.title("Learning Curve")
plt.plot(epoch_lst, train_acc_lst, label="Train")
plt.plot(epoch_lst, val_acc_lst, label="Validation")
plt.xlabel("Epoch")
plt.ylabel("Pixelwise Accuracy")
plt.legend(loc='best')
plt.show()
## train from scratch
torch.cuda.empty_cache()
TRANSFER_LEARNING = False
model = SegNet(3, 2, transfer_learning=TRANSFER_LEARNING).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
loss = nn.CrossEntropyLoss(weight=LABEL_WEIGHTS).to(DEVICE)
train(model, optimizer, loss, X_train, y_train, X_valid, y_valid)
## train from vgg weights
torch.cuda.empty_cache()
TRANSFER_LEARNING = True
model_vgginit = SegNet(3, 2, transfer_learning=TRANSFER_LEARNING).to(DEVICE)
optimizer = torch.optim.Adam(model_vgginit.parameters(), lr=LEARNING_RATE)
loss = nn.CrossEntropyLoss(weight=LABEL_WEIGHTS).to(DEVICE)
torch.cuda.empty_cache()
train(model_vgginit, optimizer, loss, X_train, y_train, X_valid, y_valid)
## load a model
model_name = r'../Epoch49_loss2.1120_trainacc97.727_valacc97.859.pth' # scratch
x, y = x.to(device), y.to(device)
out = model(x)
for (out_, y_) in zip(out, y):
out_, y_ = out_.cpu().detach().numpy(), y_.cpu().detach(
).numpy()
# print(out_.shape, y_.shape)
# print(eval_all(out_, y_))
scores.append(eval_all(out_, y_))
scores = np.array(scores)
scores = np.mean(scores, axis=0)
return scores
if __name__ == "__main__":
model = SegNet().to(device)
optimizer = torch.optim.Adam(model.parameters(),
lr=0.01,
betas=(0.9, 0.999))
train_dataset, test_dataset = load_dataset('dataset/TrainingData')
scores_avg = evaluate(test_dataset)
print(scores_avg)
for epoch in range(30):
epoch_loss = 0
for i, (x, y) in enumerate(train_dataset.batch(16)):
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
out = model(x)
transform=in_t,
target_transform=out_t,
)
v_data = Cityscapes(
'./data',
target_type=['semantic'],
split='val',
transform=in_t,
target_transform=out_t,
)
train_data = DataLoader(tr_data, batch_size=BATCH_SIZE, shuffle=True)
val_data = DataLoader(v_data, batch_size=BATCH_SIZE, shuffle=False)
# define model
model = SegNet(IN_CHANNELS, CLASSES)
if gpu:
model.to(torch.device("cuda:0"))
if monitor:
wandb.watch(model)
# optimizer and loss definition
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=LR, momentum=0.9)
# training
for epoch in range(EPOCHS):
tr_epoch_loss = 0
default=1,
help="size of the batches")
parser.add_argument("--img_size",
type=int,
default=512,
help="size of each image dimension")
parser.add_argument("--checkpoint_model",
type=str,
help="path to checkpoint model")
opt = parser.parse_args()
print(opt)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.makedirs("output", exist_ok=True)
model = SegNet().to(device)
model.load_state_dict(torch.load(opt.checkpoint_model))
model.eval()
dataloader = DataLoader(
ImageFolder(opt.image_folder, img_size=opt.img_size),
batch_size=opt.batch_size,
shuffle=False,
num_workers=4,
)
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available(
) else torch.FloatTensor
print("\nPerforming object detection:")
prev_time = time.time()
# iou = compute_iou(output, target)
# iou_meter += iou
if image_set == 'train':
optimizer.zero_grad() # Why did we do this?
loss.backward()
optimizer.step()
print('loss at epoch ', str(epoch), ' iteration ', str(i), ' is: ',
loss.data.cpu().numpy())
if __name__ == "__main__":
train_dataset = VOC('./', 'train')
test_dataset = VOC('./', 'val')
train_dataloader = data.DataLoader(train_dataset,
batch_size=8,
shuffle=False,
num_workers=4)
model = SegNet()
criterion = nn.MSELoss()
# Comment if not using a GPU
model = model.cuda()
criterion = criterion.cuda()
optimizer = torch.optim.SGD(model.parameters(), 0.01, momentum=0.9)
n_epochs = 10
for i in range(n_epochs):
train(i, train_dataloader, model, criterion, optimizer, 'train')
if args.custom:
train_dataset = ScoreDataset(list_file=train_path,
img_dir=img_dir,
mask_dir=mask_dir)
else:
train_dataset = PascalVOCDataset(list_file=train_path,
img_dir=img_dir,
mask_dir=mask_dir)
print('LOADER')
train_dataloader = DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=4)
if CUDA:
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS).cuda(GPU_ID)
class_weights = 1.0 / train_dataset.get_class_probability().cuda(
GPU_ID)
criterion = torch.nn.CrossEntropyLoss(
weight=class_weights).cuda(GPU_ID)
else:
print('MODEL')
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS)
model.init_vgg_weigts()
print('STATE_DICT')
class_weights = 1.0 / train_dataset.get_class_probability()
print('class_weights', len(class_weights))
criterion = torch.nn.CrossEntropyLoss(weight=class_weights)
writer = Writer('logs')
logger = log('')
train_dataset = PascalVOCDataset(list_file=train_path,
img_dir=train_img_path,
mask_dir=train_mask_path,
img_size=512,
is_transform=True)
train_dataloader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
model = SegNet().to(device)
class_weights = 1.0 / train_dataset.get_class_probability()
print(class_weights)
criterion = torch.nn.CrossEntropyLoss(weight=class_weights).to(device)
# start from checkpoint
if args.checkpoint:
model.load_state_dict(torch.load(args.checkpoint))
optimizer = torch.optim.SGD(model.parameters(),
lr=LEARNING_RATE,
momentum=MOMENTUM)
# training
is_better = True
prev_loss = float('inf')
from model import SegNet from data import batch_generator net = SegNet() net.fit_generator(batch_generator(), samples_per_epoch=13434, nb_epoch=5)
parser.add_argument('--batch-size', type=int, default=8)
parser.add_argument('--data_root', type=str, default='/data/pascal/VOCdevkit/VOC2012')
parser.add_argument('--val_path', type=str, default='ImageSets/Segmentation/val.txt')
parser.add_argument('--img_dir', type=str, default='JPEGImages')
parser.add_argument('--mask_dir', type=str, default='SegmentationClass')
parser.add_argument('--model_path', type=str, required=True)
args = parser.parse_args()
print(args)
data_root = args.data_root
val_dir = os.path.join(data_root, args.val_path)
img_dir = os.path.join(data_root, args.img_dir)
mask_dir = os.path.join(data_root, args.mask_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SegNet().to(device)
model.load_state_dict(torch.load(args.model_path))
val_loss, val_score, val_class_iou = validate(
model=model, val_path=val_dir, img_path=img_dir, mask_path=mask_dir, batch_size=args.batch_size
)
for k, v in val_score.items():
print(k, v)
for k, v in val_class_iou.items():
print(k, v)
print('\nval_loss:'+val_loss)
def main(args):
# set the necessary list
train_list = pd.read_csv(args.train_list, header=None)
val_list = pd.read_csv(args.val_list, header=None)
# set the necessary directories
trainimg_dir = args.trainimg_dir
trainmsk_dir = args.trainmsk_dir
valimg_dir = args.valimg_dir
valmsk_dir = args.valmsk_dir
train_gen = data_gen_small(
trainimg_dir,
trainmsk_dir,
train_list,
args.batch_size,
[args.input_shape[0], args.input_shape[1]],
args.n_labels,
)
val_gen = data_gen_small(
valimg_dir,
valmsk_dir,
val_list,
args.batch_size,
[args.input_shape[0], args.input_shape[1]],
args.n_labels,
)
model = SegNet(args.n_labels, args.kernel, args.pool_size,
args.output_mode)
model.build(input_shape=(1, 256, 256, 3))
print(model.summary())
model.compile(loss=args.loss,
optimizer=args.optimizer,
metrics=["accuracy"])
model.fit_generator(train_gen,
steps_per_epoch=args.epoch_steps,
epochs=args.n_epochs,
validation_data=val_gen,
validation_steps=args.val_steps,
callbacks=[TensorBoard(log_dir="./run")])
model.save_weights(args.save_dir + str(args.n_epochs) + ".hdf5")
print("sava weight done..")
val_dataset = VOCSegmentation(base_size=224, crop_size=224, split='val')
train_dataloader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
val_dataloader = DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
#class_weights = 1.0/calculate_weigths_labels(train_dataloader, NUM_CLASSES)
class_weights = torch.tensor([
0.5220, 0.0300, 0.0230, 0.0304, 0.0273, 0.0279, 0.0444, 0.0395, 0.0556,
0.0308, 0.0298, 0.0339, 0.0415, 0.0318, 0.0352, 0.0748, 0.0274, 0.0314,
0.0335, 0.0409, 0.0311
])
class_weights = class_weights.float().cuda()
model = SegNet(input_channels=3,
output_channels=NUM_CLASSES,
class_weights=class_weights)
model.load_vgg16_weight()
trainer = pl.Trainer(max_epochs=args.max_epochs,
gpus=args.gpus,
precision=args.precision)
#lr_finder_plot(model, trainer, train_dataloader, val_dataloader)
trainer.fit(model, train_dataloader, val_dataloader)
adversarial_path = 'data/' + args.attacks + '.pickle'
else:
adversarial_path = args.attack_path
else:
model = None
if args.model == 'UNet':
model = UNet(in_channels=n_channels, n_classes=n_classes)
elif args.model == 'SegNet':
model = SegNet(in_channels=n_channels, n_classes=n_classes)
elif args.model == 'DenseNet':
model = DenseNet(in_channels=n_channels, n_classes=n_classes)
else:
print("wrong model : must be UNet, SegNet, or DenseNet")
raise SystemExit
summary(model,
input_size=(n_channels, args.height, args.width),
device='cpu')
model.load_state_dict(torch.load(args.model_path))
adversarial_examples = DAG_Attack(model, test_dataset, args)
def train(epoch, dataloader, model, criterion, optimizer, image_set = 'train'):
loss_meter = 0
acc_meter = 0
for i, (input, target) in enumerate(dataloader):
if image_set == 'train':
input = input.requires_grad_(True).float().cuda()
else:
input = input.float().cuda()
target = target.float().cuda()
# Get the model output
output =
# Introducing the loss here. Compute the loss value
loss =
loss_meter += loss.item()
# Compute acc here
acc = compute_acc(output, target)
acc_meter += acc.item()
if image_set == 'train':
# In the next three lines:
# Zero the existing gadients
# Do a backward pass
# Update the weights
if i % 3 == 0:
print(image_set, ' loss at epoch ', str(epoch), ' iteration ', str(i), ' is: ', loss_meter / (i+1),
' and acc is: ', acc_meter / (i+1))
if __name__ == "__main__":
train_dataset = VOC('./VOCdevkit/', 'train')
val_dataset = VOC('./VOCdevkit/', 'val')
train_dataloader = data.DataLoader(
train_dataset,
batch_size = 6,
shuffle = True,
num_workers = 4)
val_dataloader = data.DataLoader(
val_dataset,
batch_size = 1,
shuffle = False,
num_workers = 1)
model = SegNet()
# criterion = nn.MSELoss()
# criterion = nn.BCELoss()
criterion = nn.BCEWithLogitsLoss()
# Comment if not using a GPU
model = model.cuda()
criterion = criterion.cuda()
# Inititialize the optimizer.
lr = 0.1
optimizer = torch.optim.Adam(model.parameters(), lr)
n_epochs = 10
for i in range(n_epochs):
train(i, train_dataloader, model, criterion, optimizer, 'train')
if i % 2 == 0:
train(i, val_dataloader, model, criterion, optimizer, 'val')
def main():
# Parse arguments.
parser = argparse.ArgumentParser()
kwargs = {
'type': int,
'default': 100,
'help': 'The number of times of learning. default: 100'
}
parser.add_argument('-e', '--epochs', **kwargs)
kwargs = {
'type': int,
'default': 10,
'help': 'The frequency of saving model. default: 10'
}
parser.add_argument('-c', '--checkpoint_interval', **kwargs)
kwargs = {
'type': int,
'default': 1,
'help': 'The number of samples contained per mini batch. default: 1'
}
parser.add_argument('-b', '--batch_size', **kwargs)
kwargs = {
'default':
False,
'action':
'store_true',
'help':
'Whether store all data to GPU. If not specified this option, use both CPU memory and GPU memory.'
}
parser.add_argument('--onmemory', **kwargs)
args = parser.parse_args()
# Prepare training data.
dataset = np.load('./temp/dataset.npz')
train_x = dataset['train_x']
train_y = dataset['train_y']
test_x = dataset['test_x']
test_y = dataset['test_y']
# Prepare tensorflow.
config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
session = tf.Session(config=config)
keras.backend.tensorflow_backend.set_session(session)
# Prepare model.
model = SegNet(shape=(360, 480, 3))
model.compile(loss='binary_crossentropy',
optimizer='adadelta',
metrics=['accuracy'])
# Training.
callbacks = []
timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
directory = f'./logs/{timestamp}/'
os.makedirs(directory, exist_ok=True)
callbacks.append(keras.callbacks.TensorBoard(log_dir=directory))
filename = 'model-{epoch:04d}.h5'
directory = f'./temp/{timestamp}/'
os.makedirs(directory, exist_ok=True)
callbacks.append(
keras.callbacks.ModelCheckpoint(filepath=f'{directory}{filename}',
monitor='val_loss',
verbose=0,
save_best_only=False,
save_weights_only=False,
mode='auto',
period=args.checkpoint_interval))
model.save_weights(f'{directory}{filename}'.format(epoch=0))
if args.onmemory:
model.fit(x=train_x,
y=train_y,
validation_data=(test_x, test_y),
epochs=args.epochs,
batch_size=args.batch_size,
class_weight='balanced',
shuffle=True,
verbose=1,
callbacks=callbacks)
else:
class Generator(keras.utils.Sequence):
def __init__(self, x, y, batch_size, shuffle):
self.x = x
self.y = y
self.batch_size = batch_size
self.indices = np.arange(len(self.x))
self.shuffle = shuffle
assert len(self.x) == len(self.y)
assert len(self.x) % self.batch_size == 0
def __getitem__(self, index):
i = index * self.batch_size
indices = self.indices[i:i + self.batch_size]
x = self.x[indices]
y = self.y[indices]
return x, y
def __len__(self):
return len(self.x) // self.batch_size
def on_epoch_end(self):
if self.shuffle:
self.indices = np.random.permutation(self.indices)
model.fit_generator(generator=Generator(train_x, train_y,
args.batch_size, True),
validation_data=Generator(test_x, test_y,
args.batch_size, False),
epochs=args.epochs,
class_weight='balanced',
shuffle=True,
verbose=1,
callbacks=callbacks)
composed = transforms.Compose([RandomCrop(256)])
file_path = "/home/huijian/exps/segnet/model/"
train_data = BuildingDataset(root_dir=train_dir, transform=composed)
test_data = BuildingDataset(root_dir=test_dir, transform=composed)
train_loader = utils.data.DataLoader(train_data,
batch_size=4,
shuffle=True)
test_loader = utils.data.DataLoader(test_data, batch_size=4, shuffle=True)
# building the net
# [64,128,256,512,512]
segnet = SegNet(input_dim=3,
output_dim=1,
features=[64, 96, 128, 256, 256])
segnet.apply(weights_normal)
print(segnet)
trainer = Trainer(net=segnet, file_path=file_path)
# restore the model
if True:
trainer.restore_model()
# begin training
if False:
print("begin training!")
trainer.train_model(train_loader=train_loader,
# Initialize Pascal VOC Dataset for validation
val_dataset = PascalVOCDataset(list_file=val_path,
img_dir=img_dir,
mask_dir=mask_dir)
# Initialize validation dataloader
val_dataloader = DataLoader(
val_dataset,
batch_size=BATCH_SIZE, # mini-batch size
shuffle=True, # Use shuffle
num_workers=4) # number of cpu
# If using cuda
if CUDA:
# Initialize SegNet model on gpu memory
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS).cuda(GPU_ID)
# Set target class values on gpu memory
class_weights = 1.0 / val_dataset.get_class_probability().cuda(GPU_ID)
# Set loss function on gpu memory
criterion = torch.nn.CrossEntropyLoss(
weight=class_weights).cuda(GPU_ID)
else:
# Initialize SegNet model on cpu memory
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS)
# Set target class values
class_weights = 1.0 / val_dataset.get_class_probability()
# Set loss function
criterion = torch.nn.CrossEntropyLoss(weight=class_weights)
sc = MinMaxScaler(feature_range=(0, 1))
# testControls = np.asarray(testControls).astype('float32')
# controls = np.asarray(controls).astype('float32')
x_train = np.asarray(images) / 255.0
x_train = np.delete(x_train, len(x_train) - 1, axis=0)
print(x_train.shape)
# sys.exit()
x_test = np.asarray(testImages) / 255.0
x_test = np.delete(x_test, len(x_test) - 1, axis=0)
y_train = sc.fit_transform(np.asarray(controls))
# Prepare model for training
model = SegNet((IMG_H, IMG_W))
model.compile(optimizer='adam', loss='mean_squared_error')
model.fit(x_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE)
if not os.path.exists('models'):
os.mkdir('models')
model.save_weights('models/res.h5', save_format='h5')
# Predict on testing dataset
predictions = model.predict(x_test)
predictions = sc.inverse_transform(predictions)
print(predictions)
from __future__ import print_function
from model import SegNet
from dataset import NUM_CLASSES
import matplotlib.pyplot as plt
import numpy as np
import torch
if __name__ == "__main__":
# RGB input
input_channels = 3
# RGB output
output_channels = NUM_CLASSES
# Model
model = SegNet(input_channels=input_channels,
output_channels=output_channels)
print(model)
img = torch.randn([4, 3, 224, 224])
# plt.imshow(np.transpose(img.numpy()[0,:,:,:],
# (1, 2, 0)))
# plt.show()
output, softmaxed_output = model(img)
# plt.imshow(np.transpose(output.detach().numpy()[0,:,:,:],
# (1, 2, 0)))
# plt.show()
def train_autoencoder(epoch_plus):
writer = SummaryWriter(log_dir='./runs_autoencoder_2')
num_epochs = 400 - epoch_plus
lr = 0.001
bta1 = 0.9
bta2 = 0.999
weight_decay = 0.001
# model = autoencoder(nchannels=3, width=172, height=600)
model = SegNet(3)
if ngpu > 1:
model = nn.DataParallel(model)
if use_gpu:
model = model.to(device, non_blocking=True)
if epoch_plus > 0:
model.load_state_dict(
torch.load('./autoencoder_models_2/autoencoder_{}.pth'.format(
epoch_plus)))
criterion = nn.MSELoss(reduction='sum')
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
betas=(bta1, bta2),
weight_decay=weight_decay)
for epoch in range(num_epochs):
degree = randint(-180, 180)
transforms = torchvision.transforms.Compose([
torchvision.transforms.CenterCrop((172, 200)),
torchvision.transforms.Resize((172, 200)),
torchvision.transforms.RandomRotation((degree, degree)),
torchvision.transforms.ToTensor()
])
dataloader = get_dataloader(data_dir,
train=True,
transform=transforms,
batch_size=batch_size)
model.train()
epoch_losses = AverageMeter()
with tqdm(total=(1000 - 1000 % batch_size)) as _tqdm:
_tqdm.set_description('epoch: {}/{}'.format(
epoch + 1 + epoch_plus, num_epochs + epoch_plus))
for data in dataloader:
gt, text = data
if use_gpu:
gt, text = gt.to(device, non_blocking=True), text.to(
device, non_blocking=True)
predicted = model(text)
# loss = criterion_bce(predicted, gt) + criterion_dice(predicted, gt)
loss = criterion(
predicted, gt - text
) # predicts extracted text in white, all others in black
epoch_losses.update(loss.item(), len(gt))
optimizer.zero_grad()
loss.backward()
optimizer.step()
_tqdm.set_postfix(loss='{:.6f}'.format(epoch_losses.avg))
_tqdm.update(len(gt))
save_path = './autoencoder_models_2'
if not os.path.exists(save_path):
os.mkdir(save_path)
gt_text = gt - text
predicted_mask = text + predicted
torch.save(
model.state_dict(),
os.path.join(save_path,
'autoencoder_{}.pth'.format(epoch + 1 + epoch_plus)))
writer.add_scalar('Loss', epoch_losses.avg, epoch + 1 + epoch_plus)
writer.add_image('text/text_image_{}'.format(epoch + 1 + epoch_plus),
text[0].squeeze(), epoch + 1 + epoch_plus)
writer.add_image('gt/gt_image_{}'.format(epoch + 1 + epoch_plus),
gt[0].squeeze(), epoch + 1 + epoch_plus)
writer.add_image('gt_text/gt_image_{}'.format(epoch + 1 + epoch_plus),
gt_text[0].squeeze(), epoch + 1 + epoch_plus)
writer.add_image(
'predicted/predicted_image_{}'.format(epoch + 1 + epoch_plus),
predicted_mask[0].squeeze(), epoch + 1 + epoch_plus)
writer.add_image(
'predicted_text/predicted_image_{}'.format(epoch + 1 + epoch_plus),
predicted[0].squeeze(), epoch + 1 + epoch_plus)
writer.close()