def load_images(folder, batch_size, k):
train_filename = '{}/train.large.input.npy'.format(folder)
valid_filename = '{}/val.input.npy'.format(folder)
test_filename = '{}/test.input.npy'.format(folder)
noise_filename = '{}/noise.input.npy'.format(folder)
train_metadata = pickle.load(open('{}/train.large.metadata.p'.format(folder), 'rb'))
valid_metadata = pickle.load(open('{}/val.metadata.p'.format(folder), 'rb'))
test_metadata = pickle.load(open('{}/test.metadata.p'.format(folder), 'rb'))
noise_metadata = pickle.load(open('{}/noise.metadata.p'.format(folder), 'rb'))
train_dataset = ImageDataset(train_filename)
valid_dataset = ImageDataset(valid_filename, mean=train_dataset.mean, std=train_dataset.std) # All features are normalized with mean and std
test_dataset = ImageDataset(test_filename, mean=train_dataset.mean, std=train_dataset.std)
noise_dataset = ImageDataset(noise_filename, mean=train_dataset.mean, std=train_dataset.std)
train_data = DataLoader(train_dataset, num_workers=1, pin_memory=True,
batch_sampler=BatchSampler(ImagesSampler(train_dataset, train_metadata, k, shuffle=True), batch_size=batch_size, drop_last=False))
valid_data = DataLoader(valid_dataset, num_workers=1, pin_memory=True,
batch_sampler=BatchSampler(ImagesSampler(valid_dataset, valid_metadata, k, shuffle=False), batch_size=batch_size, drop_last=False))
test_data = DataLoader(test_dataset, num_workers=1, pin_memory=True,
batch_sampler=BatchSampler(ImagesSampler(test_dataset, test_metadata, k, shuffle=False), batch_size=batch_size, drop_last=False))
noise_data = DataLoader(noise_dataset, num_workers=1, pin_memory=True,
batch_sampler=BatchSampler(ImagesSampler(noise_dataset, noise_metadata, k, shuffle=False), batch_size=batch_size, drop_last=False))
return train_data, valid_data, test_data, noise_data
def load_data(input, batch=4):
camera_mat = camera()
train_dir = PATH + 'train_images/'
train, validate = train_test_split(input, test_size=0.01, random_state=13)
train_data = ImageDataset(train, train_dir, camera_mat)
validate_data = ImageDataset(validate, train_dir, camera_mat)
train_loader = DataLoader(dataset=train_data,
batch_size=batch,
shuffle=True,
num_workers=2)
validate_loader = DataLoader(dataset=validate_data,
batch_size=batch,
shuffle=False,
num_workers=0)
return train_loader, validate_loader, validate_data, validate
def __init__(self, root: str, test_dir: str, dst_img_size: Tuple[int,
int]):
self._search_for_files(os.path.join(root, test_dir))
self.cache = ImageCache(force_rgb=True, dst_img_size=dst_img_size)
# loading labels
self.labels = ImageDataset.load_labels(root)
if self.labels is None:
print("Failed to load labels. Data root is " + root)
return
# checking for extra classes
labels_set = set(self.labels)
filtered_files = []
for filename in self.image_files:
label = self.image_labels[filename]
if label not in labels_set:
# raise Exception('test label not in training set: ' + label)
pass
else:
filtered_files.append(filename)
if len(filtered_files) != len(self.image_files):
print(
'Test files filtered to conform to training labels. %d files were excluded'
% (len(self.image_files) - len(filtered_files)))
self.image_files = filtered_files
# indexing labels
self.label_to_index = dict()
for i, label in enumerate(self.labels):
self.label_to_index[label] = i
def _build_automated_dataset(self,
name,
file_list,
img_dir,
test=True,
transform=None,
batch_size=None,
num_workers=1,
shuffle=False):
if transform is None:
transform = self.test_transform
if batch_size is None:
batch_size = self.test_batch_size
data = ImageDataset(csv_file=file_list,
img_dir=img_dir,
transform=transform,
test=test)
loader = DataLoader(data,
batch_size=batch_size,
shuffle=shuffle,
pin_memory=True,
num_workers=num_workers)
automated_dataset = AutomatedDataset(name, loader)
return automated_dataset
def predict_from_dir(args):
images = os.listdir(args.dir)
if len(images) == 0:
print('Cannot find any images in ' + args.dir)
return
print('%d images to classify found' % len(images))
use_cuda = torch.cuda.is_available() and args.gpu.lower() == 'true'
device = torch.device("cuda" if use_cuda else "cpu")
print("Using GPU" if use_cuda else "Using CPU")
# loading labels
labels = ImageDataset.load_labels(args.data_root)
if labels is None:
print("Failed to load labels. Data root is " + args.data_root)
return
num_classes = len(labels)
print("Number of recognized classes = %d" % num_classes)
# loading model
print("loading model...")
model, model_name = load_model(args.model,
device,
num_classes,
inference=True)
target_size = (224, 224)
# loading snapshot
snapshot_file = os.path.join("snapshots", model_name,
"snapshot_%d.pt" % args.snapshot)
if not os.path.exists(snapshot_file):
print("Snapshot file does not exists: " + snapshot_file)
return
model.load_state_dict(torch.load(snapshot_file, 'cpu'))
for image_name in images:
img = cv2.imread(os.path.join(args.dir, image_name))
if img is None:
print('Failed to load ' + image_name)
continue
img = cv2.resize(img, target_size, interpolation=cv2.INTER_AREA)
# getting required number of frames for prediction
images_to_predict = [img]
probs = predict_with_tta(args.tta, images_to_predict, model,
target_size)
sorted_indexes = np.argsort(probs)
for t in range(5):
label_index = sorted_indexes[-t - 1]
label_prob = probs[label_index]
label_name = labels[label_index]
s = "%s: %1.2f%% %s" % (image_name, label_prob * 100.0, label_name)
print(s)
def compute_scattering(dir_images, dir_to_save, scattering, batch_size):
dir_to_save.mkdir()
dataset = ImageDataset(dir_images)
dataloader = DataLoader(dataset,
batch_size,
pin_memory=True,
num_workers=1)
for idx_batch, current_batch in enumerate(tqdm(dataloader)):
images = current_batch.float().cuda()
if scattering is not None:
s_images = scattering(images).cpu().numpy()
s_images = np.reshape(
s_images,
(batch_size, -1, s_images.shape[-1], s_images.shape[-1]))
else:
s_images = images.cpu().numpy()
for idx_local in range(batch_size):
idx_global = idx_local + idx_batch * batch_size
filename = dir_to_save / '{}.npy'.format(idx_global)
temp = s_images[idx_local]
np.save(filename, temp)
def get_train_val_embeddings(triplet_model, train_tracks, val_tracks):
train_images = get_images_from_tracks(train_tracks)
val_images = get_images_from_tracks(val_tracks)
split_index = len(train_images)
total_images = train_images + val_images
total_images_dataset = ImageDataset(total_images, Path("image_train"),
(256, 256))
total_images_dataloader = DataLoader(total_images_dataset,
num_workers=4,
batch_size=64,
shuffle=False)
total_embeddings = []
with torch.no_grad():
triplet_model = triplet_model.eval()
for image_batch in tqdm(total_images_dataloader):
image_batch = image_batch.cuda()
batch_embeddings = triplet_model.getEmbedding(image_batch)
total_embeddings.append(batch_embeddings.cpu())
total_embeddings = torch.cat(total_embeddings, dim=0).numpy()
val_embeddings = total_embeddings[split_index:, :]
return total_embeddings, val_embeddings
def main(args):
batch_size = 32
snapshot_iters = 500
test_iters = 100
snapshot_dir = "snapshots"
image_size = (224, 224)
learning_rate = 0.01
momentum = 0.9
debug_dir = "debug"
debug_number = 300 if args.debug.lower() == 'true' else 0
freeze_encoder = args.freeze_encoder.lower() == 'true'
set_num_threads(1)
use_cuda = torch.cuda.is_available() and args.gpu.lower() == 'true'
device = torch.device("cuda" if use_cuda else "cpu")
print("Using GPU" if use_cuda else "Using CPU")
# create data sets
train_foregrounds = ImageDataset(os.path.join(args.data_root, "train.txt"), "sorted")
val_foregrounds = ImageDataset(os.path.join(args.data_root, "val.txt"), "sorted")
num_classes = len(train_foregrounds.labels)
backgrounds = ImageDataset(os.path.join(args.data_root, "backgrounds"), force_rgb=True)
# initialize test set
test_set = TestDataset(args.data_root, "test", dst_img_size=image_size)
print("Number of classes = %d" % num_classes)
print("Number of train foregrounds = %d" % len(train_foregrounds))
print("Number of validation foregrounds = %d" % len(val_foregrounds))
print("Number of test images = %d" % len(test_set))
print("Number of backgrounds = %d" % len(backgrounds))
if len(train_foregrounds) == 0 or len(val_foregrounds) == 0 or len(backgrounds) == 0:
raise Exception("One of datasets is empty")
train_dataset = RenderingDataset(backgrounds, train_foregrounds, image_size)
test_dataset = RenderingDataset(backgrounds, val_foregrounds, image_size)
# create data loaders
kwargs = {'num_workers': 6, 'pin_memory': True} if use_cuda else {}
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, **kwargs)
val_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True, **kwargs)
# creating model
model, model_name = load_model(args.model, device, num_classes, inference=False, freeze_encoder=freeze_encoder)
# preparing snapshot dir
if not os.path.exists(snapshot_dir):
os.mkdir(snapshot_dir)
# preparing model dir
model_dir = os.path.join(snapshot_dir, model_name)
if not os.path.exists(model_dir):
os.mkdir(model_dir)
iteration = 0
if args.snapshot is not None:
iteration = args.snapshot
snapshot_file = get_snapshot_file_name(iteration, model_dir)
print("loading " + snapshot_file)
model.load_state_dict(torch.load(snapshot_file))
print('Starting from iteration %d' % iteration)
# creating optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)
# creating logs
train_losses = {}
train_accuracies = {}
val_losses = {}
val_accuracies = {}
test_accuracies = {}
# creating debug dumper
debug_dumper = DebugDumper(debug_dir, debug_number)
# creating log files
train_log_file = os.path.join(model_dir, 'train.log')
val_log_file = os.path.join(model_dir, 'val.log')
test_log_file = os.path.join(model_dir, 'test.log')
while iteration < args.num_iterations:
# train
print("training...")
iteration = train(train_loader, model, snapshot_iters, iteration, device, train_losses, train_accuracies,
optimizer, debug_dumper, train_log_file)
# dumping snapshot
snapshot_file = get_snapshot_file_name(iteration, model_dir)
print("dumping snapshot: " + snapshot_file)
torch.save(model.state_dict(), snapshot_file)
# validation
print("validation...")
train(val_loader, model, test_iters, iteration, device, val_losses, val_accuracies,
log_file_name=val_log_file)
# test
print("test...")
test(test_set, model, iteration, device, test_accuracies, test_log_file)
# visualizing training progress
plot_histograms(model, model_dir)
plot_train_curves(train_losses, val_losses, "TrainCurves", model_dir)
plot_train_curves(train_accuracies, val_accuracies, "Validation Accuracy", model_dir, logy=False)
plot_train_curves(val_accuracies, test_accuracies, "Test Accuracy", model_dir, logy=False)
if args.v >= 1: print 'Loading datasets... ',
t_load = transforms.Compose([transforms.Pad(4)])
t_train = transforms.Compose(
[transforms.RandomRotation(5),
transforms.Resize((64, 64))])
dataset_folder = '/public/bbeyers/CSC249_project/mnist/training'
dataset_csv = '/public/bbeyers/CSC249_project/mnist/training.csv'
if args.quick_load:
train_dataset = ImageDataset(dataset_folder,
dataset_csv,
transform=t_load,
train_transform=t_train,
stop_after=args.batch_size)
else:
train_dataset = ImageDataset(dataset_folder,
dataset_csv,
transform=t_load,
train_transform=t_train)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
if args.v >= 1: print 'Done!'
generator = BasicGenerator(z_dim=100, use_bias=True)
ssim = SSIM(inference, target).item()
self.log(f'{step_name}_ssim', ssim, on_epoch=True, prog_bar=True)
def validation_step(self, batch, batch_idx):
self.evaluation_step(batch, batch_idx, "val")
def test_step(self, batch, batch_idx):
self.evaluation_step(batch, batch_idx, "test")
if __name__ == '__main__':
batch_size = 2
training_set = ImageDataset([
"Dataset/classroom/inputs", "Dataset/living-room/inputs",
"Dataset/sponza/inputs", "Dataset/sponza-(glossy)/inputs",
"Dataset/sponza-(moving-light)/inputs"
])
validation_set = ImageDataset(["Dataset/san-miguel/inputs"])
validation_data_loader = DataLoader(validation_set,
batch_size=batch_size,
num_workers=8)
logger = TensorBoardLogger('tensorboard',
name='GlobalIlluminationFiltering')
log_dir = logger.log_dir
model = GlobalIlluminationFiltering(log_dir)
trainer = pl.Trainer(max_epochs=16,
gpus=1,
parser = argparse.ArgumentParser(description='classification')
# 图片数据的根目录(Root catalog of images)
parser.add_argument('--data-dir', type=str, default='images')
parser.add_argument('--batch-size', type=int, default=16)
parser.add_argument('--num-epochs', type=int, default=150)
parser.add_argument('--lr', type=float, default=0.045)
parser.add_argument('--num-workers', type=int, default=4)
parser.add_argument('--print-freq', type=int, default=1)
parser.add_argument('--save-epoch-freq', type=int, default=1)
parser.add_argument('--save-path', type=str, default='output')
parser.add_argument('--resume', type=str, default='', help='For training from one checkpoint')
parser.add_argument('--start-epoch', type=int, default=0, help='Corresponding to the epoch of resume')
args = parser.parse_args()
# read data
dataloders, dataset_sizes, class_names = ImageDataset(args)
with open('class_names.json', 'w') as f:
json.dump(class_names, f)
# use gpu or not
use_gpu = torch.cuda.is_available()
print("use_gpu:{}".format(use_gpu))
# get model
model = SimpleNet()
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
model.load_state_dict(torch.load(args.resume))
def video_capture(args):
use_cuda = torch.cuda.is_available() and args.gpu.lower() == 'true'
device = torch.device("cuda" if use_cuda else "cpu")
print("Using GPU" if use_cuda else "Using CPU")
# loading train foregrounds
train_foregrounds = ImageDataset(os.path.join(args.data_root, "train.txt"),
"sorted")
# loading model
print("loading model...")
model = modelA(num_classes, False).to(device)
model_name = type(model).__name__
model, labels_map, int_to_labels_map = load_model(args.data_root,
args.model,
args.snapshot)
print("start video capture...")
cap = cv2.VideoCapture(args.camera)
# Check if camera opened successfully
if not cap.isOpened():
print("Error opening video stream or file")
return
font = cv2.FONT_HERSHEY_SIMPLEX
past_frames = []
while cap.isOpened():
# Capture frame-by-frame
ret, frame = cap.read()
if ret:
past_frames.append(clip(frame))
if len(past_frames) > args.rtta:
past_frames.pop(0)
frame_to_show = cv2.resize(past_frames[-1],
(frame.shape[1], frame.shape[0]),
interpolation=cv2.INTER_AREA)
probs = predict_with_tta(args.tta, True, past_frames, model)[0]
sorted_indexes = np.argsort(probs)
for t in range(5):
label_index = sorted_indexes[-t - 1]
label_prob = probs[label_index]
label_name = int_to_labels_map[label_index]
s = "%1.2f%% %s" % (label_prob * 100.0, label_name)
cv2.putText(frame_to_show, s, (0, (t + 1) * 32), font, 1,
(0, 0, 180), 2, cv2.LINE_AA)
# Display the resulting frame
cv2.imshow('Frame', frame_to_show)
# Press Q on keyboard to exit
if cv2.waitKey(25) & 0xFF == ord('q'):
break
# Break the loop
else:
break
# When everything done, release the video capture object
cap.release()
# Closes all the frames
cv2.destroyAllWindows()
import os
from ImageDataset import ImageDataset
from torch.utils.data import DataLoader
from functions import *
import numpy as np
import glob
device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
paramPath = 'storedModels/SmoothL1Loss-sgd-bestParam.pth'
model = VggBn_Net()
model = model.to(device)
model.load_state_dict(torch.load(paramPath, map_location=device))
model.eval()
image_datasets = {x: ImageDataset(dataType=x) for x in ['train', 'test']}
train_len = len(image_datasets['train'])
test_len = len(image_datasets['test'])
dataloaders_dict = {}
dataloaders_dict['train'] = DataLoader(image_datasets['train'],
batch_size=8,
shuffle=False,
num_workers=4)
dataloaders_dict['test'] = DataLoader(image_datasets['test'],
batch_size=8,
shuffle=False,
num_workers=4)
phase = ['train', 'test']
'''
def video_capture(args):
use_cuda = torch.cuda.is_available() and args.gpu.lower() == 'true'
device = torch.device("cuda" if use_cuda else "cpu")
print("Using GPU" if use_cuda else "Using CPU")
# loading labels
labels = ImageDataset.load_labels(args.data_root)
if labels is None:
print("Failed to load labels. Data root is " + args.data_root)
return
num_classes = len(labels)
print("Number of recognized classes = %d" % num_classes)
if not os.path.exists(args.dump_dir):
print('Creating dump directory: ' + args.dump_dir)
os.mkdir(args.dump_dir)
# loading model
print("loading model...")
model, model_name = load_model(args.model,
device,
num_classes,
inference=True)
target_size = (224, 224)
# loading snapshot
snapshot_file = os.path.join("snapshots", model_name,
"snapshot_%d.pt" % args.snapshot)
if not os.path.exists(snapshot_file):
print("Snapshot file does not exists: " + snapshot_file)
return
model.load_state_dict(torch.load(snapshot_file, 'cpu'))
print("start video capture...")
cap = cv2.VideoCapture(args.camera)
# Check if camera opened successfully
if not cap.isOpened():
print("Error opening video stream or file")
return
font = cv2.FONT_HERSHEY_SIMPLEX
past_frames = []
dump_image_counter = 1
while cap.isOpened():
# Capture frame-by-frame
ret, frame = cap.read()
if ret:
past_frames.append(clip(frame))
if len(past_frames) > max(1, args.rtta):
past_frames.pop(0)
frame_to_show = cv2.resize(past_frames[-1],
(frame.shape[1], frame.shape[0]),
interpolation=cv2.INTER_AREA)
# getting required number of frames for prediction
images_to_predict = []
if args.rtta < 1:
images_to_predict.append(past_frames[-1])
else:
images_to_predict = past_frames[0:min(args.
rtta, len(past_frames))]
probs = predict_with_tta(args.tta, images_to_predict, model,
target_size)
sorted_indexes = np.argsort(probs)
frame_to_show_saved = frame_to_show.copy()
for t in range(5):
label_index = sorted_indexes[-t - 1]
label_prob = probs[label_index]
label_name = labels[label_index]
s = "%1.2f%% %s" % (label_prob * 100.0, label_name)
cv2.putText(frame_to_show, s, (0, (t + 1) * 32), font, 1,
(0, 0, 180), 2, cv2.LINE_AA)
# Display the resulting frame
cv2.imshow('Frame', frame_to_show)
# Press Q on keyboard to exit
key = cv2.waitKey(25) & 0xFF
if key == ord('q'):
break
if key == ord('p'):
top_label = labels[sorted_indexes[-1]]
label_dir = os.path.join(args.dump_dir, top_label)
if not os.path.exists(label_dir):
os.mkdir(label_dir)
image_file_name = os.path.join(
label_dir, 'img_%d.jpg' % dump_image_counter)
cv2.imwrite(image_file_name, frame_to_show_saved)
dump_image_counter += 1
# Break the loop
else:
break
# When everything done, release the video capture object
cap.release()
# Closes all the frames
cv2.destroyAllWindows()
def __init__(self, config):
torch.manual_seed(config.seed)
self.config = config
self.train_transform = transforms.Compose([
#transforms.RandomRotation(3),
AdaptiveResize(512),
transforms.RandomCrop(config.image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
self.test_transform = transforms.Compose([
AdaptiveResize(768),
transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
self.train_batch_size = config.batch_size
self.test_batch_size = 1
self.ranking = config.ranking
self.train_data = ImageDataset(csv_file=os.path.join(
config.trainset, 'splits2', str(config.split), config.train_txt),
img_dir=config.trainset,
transform=self.train_transform,
test=(not config.ranking))
self.train_loader = DataLoader(self.train_data,
batch_size=self.train_batch_size,
shuffle=True,
pin_memory=True,
num_workers=12)
# testing set configuration
self.live_data = ImageDataset(csv_file=os.path.join(
config.live_set, 'splits2', str(config.split), 'live_test.txt'),
img_dir=config.live_set,
transform=self.test_transform,
test=True)
self.live_loader = DataLoader(self.live_data,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=True,
num_workers=1)
self.csiq_data = ImageDataset(csv_file=os.path.join(
config.csiq_set, 'splits2', str(config.split), 'csiq_test.txt'),
img_dir=config.csiq_set,
transform=self.test_transform,
test=True)
self.csiq_loader = DataLoader(self.csiq_data,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=True,
num_workers=1)
self.tid2013_data = ImageDataset(csv_file=os.path.join(
config.tid2013_set, 'splits2', str(config.split), 'tid_test.txt'),
img_dir=config.tid2013_set,
transform=self.test_transform,
test=True)
self.tid2013_loader = DataLoader(self.tid2013_data,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=True,
num_workers=1)
self.kadid10k_data = ImageDataset(csv_file=os.path.join(
config.kadid10k_set, 'splits2', str(config.split),
'kadid10k_test.txt'),
img_dir=config.kadid10k_set,
transform=self.test_transform,
test=True)
self.kadid10k_loader = DataLoader(self.kadid10k_data,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=True,
num_workers=1)
self.bid_data = ImageDataset(csv_file=os.path.join(
config.bid_set, 'splits2', str(config.split), 'bid_test.txt'),
img_dir=config.bid_set,
transform=self.test_transform,
test=True)
self.bid_loader = DataLoader(self.bid_data,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=True,
num_workers=1)
self.clive_data = ImageDataset(csv_file=os.path.join(
config.clive_set, 'splits2', str(config.split), 'clive_test.txt'),
img_dir=config.clive_set,
transform=self.test_transform,
test=True)
self.clive_loader = DataLoader(self.clive_data,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=True,
num_workers=1)
self.koniq10k_data = ImageDataset(csv_file=os.path.join(
config.koniq10k_set, 'splits2', str(config.split),
'koniq10k_test.txt'),
img_dir=config.koniq10k_set,
transform=self.test_transform,
test=True)
self.koniq10k_loader = DataLoader(self.koniq10k_data,
batch_size=self.test_batch_size,
shuffle=False,
pin_memory=True,
num_workers=1)
self.device = torch.device(
"cuda" if torch.cuda.is_available() and config.use_cuda else "cpu")
# initialize the model
#self.model = E2EUIQA()
if config.network == 'basecnn':
self.model = BaseCNN(config)
self.model = nn.DataParallel(self.model, device_ids=[0])
elif config.network == 'dbcnn':
self.model = DBCNN(config)
self.model = nn.DataParallel(self.model).cuda()
else:
raise NotImplementedError(
"Not supported network, need to be added!")
self.model.to(self.device)
self.model_name = type(self.model).__name__
print(self.model)
# loss function
if config.ranking:
if config.fidelity:
self.loss_fn = Fidelity_Loss()
else:
self.loss_fn = nn.BCEWithLogitsLoss()
else:
self.loss_fn = nn.MSELoss()
self.loss_fn.to(self.device)
if self.config.std_modeling:
self.std_loss_fn = nn.MarginRankingLoss(margin=self.config.margin)
self.std_loss_fn.to(self.device)
self.initial_lr = config.lr
if self.initial_lr is None:
lr = 0.0005
else:
lr = self.initial_lr
#self.optimizer = optim.Adam([{'params': self.model.parameters(), 'lr': lr},
# {'params': self.sensitivity, 'lr': 1e-3},
# {'params': self.specificity, 'lr': 1e-3}]
# )
#self.optimizer = optim.Adam([{'params': self.model.backbone.parameters(), 'lr': lr},
# {'params': self.model.fc.parameters(), 'lr': 10 * lr}],
# lr=lr, weight_decay=5e-4
# )
#self.optimizer = optim.SGD([{'params': self.model.backbone.parameters(), 'lr': lr},
# {'params': self.model.fc.parameters(), 'lr': lr}],
# lr=lr, weight_decay=5e-4, momentum=0.9
# )
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=lr,
weight_decay=5e-4)
# some states
self.start_epoch = 0
self.start_step = 0
self.train_loss = []
self.train_std_loss = []
self.test_results_srcc = {
'live': [],
'csiq': [],
'tid2013': [],
'kadid10k': [],
'bid': [],
'clive': [],
'koniq10k': []
}
self.test_results_plcc = {
'live': [],
'csiq': [],
'tid2013': [],
'kadid10k': [],
'bid': [],
'clive': [],
'koniq10k': []
}
self.ckpt_path = config.ckpt_path
self.max_epochs = config.max_epochs
self.epochs_per_eval = config.epochs_per_eval
self.epochs_per_save = config.epochs_per_save
# try load the model
if config.resume or not config.train:
if config.ckpt:
ckpt = os.path.join(config.ckpt_path, config.ckpt)
else:
ckpt = self._get_latest_checkpoint(path=config.ckpt_path)
self._load_checkpoint(ckpt=ckpt)
self.scheduler = lr_scheduler.StepLR(self.optimizer,
last_epoch=self.start_epoch - 1,
step_size=config.decay_interval,
gamma=config.decay_ratio)
import glob
import random
import os
import numpy as np
import torch
from torch.utils.data import Dataset
from PIL import Image
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torchvision.datasets as dset
from ImageDataset import ImageDataset
hr_transforms = transforms.Compose([
# transforms.Resize(256),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
lr_transforms = transforms.Compose([
# transforms.Resize(64),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataloader = DataLoader(ImageDataset(lr_transforms=lr_transforms,
hr_transforms=hr_transforms),
batch_size=256,
shuffle=True,
num_workers=4)
def __init__(self, options, path):
"""Prepare the network, criterion, solver, and data.
Args:
options, dict: Hyperparameters.
"""
print('Prepare the network and data.')
self._options = options
self._path = path
self._epoch = 0
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
# Network.
network = SCNN()
#weight_init(network)
network = network.to(self.device)
# self._net = network.cuda()
self._net = network
#self._net = torch.nn.DataParallel(network)
logspaced_LR = np.logspace(-1, -4, self._options['epochs'])
# Load the model from disk.
checkpoints_list = os.listdir(self._path['model'])
if len(checkpoints_list) != 0:
self._net.load_state_dict(
torch.load(
os.path.join(
self._path['model'],
'%s%s%s' % ('net_params',
str(len(checkpoints_list) - 1), '.pkl'))))
self._epoch = len(checkpoints_list)
self._options['base_lr'] = logspaced_LR[len(checkpoints_list)]
# self._net.load_state_dict(torch.load(self._path['model']))
print(self._net)
# Criterion.
self._criterion = torch.nn.CrossEntropyLoss().cuda()
# Solver.
self._solver = torch.optim.SGD(
self._net.parameters(),
lr=self._options['base_lr'],
momentum=0.9,
weight_decay=self._options['weight_decay'])
# self._solver = torch.optim.Adam(
# self._net.parameters(), lr=self._options['base_lr'],
# weight_decay=self._options['weight_decay'])
lambda1 = lambda epoch: logspaced_LR[epoch]
self._scheduler = torch.optim.lr_scheduler.LambdaLR(self._solver,
lr_lambda=lambda1)
self.train_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=256), # Let smaller edge match
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.RandomCrop(size=224),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
self.test_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=256),
torchvision.transforms.CenterCrop(size=224),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
self.train_data = ImageDataset(
csv_file=os.path.join(path['kadis'], 'train.txt'),
img_dir=os.path.join(path['kadis'], 'dist_imgs'),
transform=self.train_transforms,
test=False)
self._train_loader = torch.utils.data.DataLoader(
self.train_data,
batch_size=self._options['batch_size'],
shuffle=True,
num_workers=1,
pin_memory=True)
self.test_data = ImageDataset(
csv_file=os.path.join(path['kadis'], 'test.txt'),
img_dir=os.path.join(path['kadis'], 'dist_imgs'),
transform=self.test_transforms,
test=True)
self._test_loader = torch.utils.data.DataLoader(
self.test_data,
batch_size=self._options['batch_size'],
shuffle=False,
num_workers=1,
pin_memory=True)
model = Alex_Net()
initial_lr = 1e-3
model = model.to(device)
#choose different update rule
params_to_update = model.parameters()
if (args['update_rule'] == 'Adam'):
optimizer = torch.optim.Adam(params_to_update, lr=initial_lr)
else:
optimizer = torch.optim.SGD(params_to_update, lr=initial_lr, momentum=0.9)
#get train and val dataloaders
#and choose whether data_augmentation
image_datasets = {
x: ImageDataset(dataType=x, dataAug=(args['data_aug'] == 1))
for x in ['train', 'val']
}
dataloaders_dict = {
x: DataLoader(image_datasets[x],
batch_size=args['batchsize'],
shuffle=True,
num_workers=args['num_workers'])
for x in ['train', 'val']
}
#choose different loss function
if (args['loss_type'] == 'L1'):
criterions = [nn.CrossEntropyLoss(), nn.L1Loss()]
elif (args['loss_type'] == 'SmoothL1Loss'):
criterions = [nn.CrossEntropyLoss(), nn.SmoothL1Loss()]
# Structural similarity index measure.
# Result is in the range [-1, 1], where -1 is completely dissimilar and 1 is perfectly similar.
# http://www.cns.nyu.edu/pub/lcv/wang03-reprint.pdf
# https://en.wikipedia.org/wiki/Structural_similarity
# https://medium.com/srm-mic/all-about-structural-similarity-index-ssim-theory-code-in-pytorch-6551b455541e
def SSIM(noisy_tensor, reference_tensor, filter_std_dev=1.5):
return SSIM_pr_pixel(noisy_tensor, reference_tensor, filter_std_dev).mean()
if __name__ == '__main__':
from ImageDataset import ImageDataset
from Visualize import show_HDR_tensor
training_set = ImageDataset(["Dataset/san-miguel/inputs"])
(light_tensor, albedo_tensor, _, _), reference_tensor = training_set[0]
color_tensor = light_tensor * albedo_tensor
identity_psnr = PSNR(reference_tensor, reference_tensor).item()
print(f"PSNR of same images: {identity_psnr:.6f}")
noisy_psnr = PSNR(color_tensor, reference_tensor).item()
print(f"PSNR of noisy and reference images: {noisy_psnr:.6f}")
identity_ssim = SSIM(reference_tensor, reference_tensor).item()
print(f"SSIM of same images: {identity_ssim:.6f}")
noisy_ssim = SSIM(color_tensor, reference_tensor).item()
print(f"SSIM of noisy and reference images: {noisy_ssim:.6f}")
def main(args):
batch_size = 32
snapshot_iters = 500
test_iters = 100
snapshot_dir = "snapshots"
image_size = (224, 224)
learning_rate = 0.01
debug_dir = "debug"
debug_number = 0
use_cuda = torch.cuda.is_available() and args.gpu.lower() == 'true'
device = torch.device("cuda" if use_cuda else "cpu")
print("Using GPU" if use_cuda else "Using CPU")
# create data sets
train_foregrounds = ImageDataset(os.path.join(args.data_root, "train.txt"),
"sorted")
test_foregrounds = ImageDataset(os.path.join(args.data_root, "val.txt"),
"sorted")
num_classes = len(train_foregrounds.labels)
backgrounds = ImageDataset(os.path.join(args.data_root, "backgrounds"),
forceRgb=True)
print("Number of classes = %d" % num_classes)
print("Number of train foregrounds = %d" % len(train_foregrounds))
print("Number of test foregrounds = %d" % len(test_foregrounds))
print("Number of backgrounds = %d" % len(backgrounds))
if len(train_foregrounds) == 0 or len(test_foregrounds) == 0 or len(
backgrounds) == 0:
raise Exception("One of datasets is empty")
train_dataset = RenderingDataset(backgrounds, train_foregrounds,
image_size)
test_dataset = RenderingDataset(backgrounds, test_foregrounds, image_size)
# create data loaders
kwargs = {'num_workers': 6, 'pin_memory': False} if use_cuda else {}
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
**kwargs)
# creating model
model = modelA(num_classes, False).to(device)
model_name = type(model).__name__
# preparing snapshot dir
if not os.path.exists(snapshot_dir):
os.mkdir(snapshot_dir)
# preparing model dir
model_dir = os.path.join(snapshot_dir, model_name)
if not os.path.exists(model_dir):
os.mkdir(model_dir)
iteration = 0
if args.snapshot is not None:
iteration = args.snapshot
snapshot_file = get_snapshot_file_name(iteration, model_dir)
print("loading " + snapshot_file)
model.load_state_dict(torch.load(snapshot_file))
print('Starting from iteration %d' % iteration)
# creating optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
# creating logs
train_losses = {}
train_accuracies = {}
test_losses = {}
test_accuracies = {}
# creating debug dumper
debug_dumper = DebugDumper(debug_dir, debug_number)
while iteration < args.num_iterations:
# train
print("training...")
iteration = train(train_loader, model, snapshot_iters, iteration,
device, train_losses, train_accuracies, optimizer,
debug_dumper)
# dumping snapshot
snapshot_file = get_snapshot_file_name(iteration, model_dir)
print("dumping snapshot: " + snapshot_file)
torch.save(model.state_dict(), snapshot_file)
# test
print("testing...")
train(test_loader, model, test_iters, iteration, device, test_losses,
test_accuracies)
# visualizing training progress
plot_histograms(model, model_dir)
plot_train_curves(train_losses, test_losses, "TrainCurves", model_dir)
plot_train_curves(train_accuracies, test_accuracies, "Accuracy",
model_dir)
parser.add_argument("--latent_dim", type=int, default=100, help="dimensionality of the latent space")
parser.add_argument("--img_size", type=int, default=28, help="size of each image dimension")
parser.add_argument("--channels", type=int, default=3, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=400, help="interval betwen image samples")
opt = parser.parse_args()
print(opt)
img_shape = (opt.channels, opt.img_size, opt.img_size)
cuda = True if torch.cuda.is_available() else False
# Dataset and Dataloader
transform = transforms.Compose([
transforms.Resize(128)
])
dataset = ImageDataset('../data/celeba', transform=transform)
batch_size = 16
data_loader = loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
# Model
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__
def block(in_feat, out_feat, normalize = True):
layers = [nn.Linear(in_feat, out_feat)]
if normalize:
layers.append(nn.BatchNorm1d(out_feat, 0.8))
# ### Load DataSet
# In[9]:
# 256 x 256
transform_train = [
transforms.Resize(int(opt.size * 1.12), Image.BICUBIC),
transforms.RandomCrop(opt.size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]
#transform_test = [transforms.ToTensor(),
# transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
trainset = ImageDataset(opt.dataroot,
transforms_=transform_train,
mode='train')
#testset = ImageDataset(opt.dataroot, transforms_=transform_test, mode='test')
train_loader = DataLoader(trainset, batch_size=opt.batchSize, shuffle=True)
#test_loader = DataLoader(testset, batch_size=opt.batchSize, shuffle=True)
print("Number of training data:", len(trainset))
#print(len(testset))
# ### Declaration
# In[10]:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CycleGAN(opt).to(device)
#process path for log, output plot and model
for b in range(0, batch_size):
loss = model.loss_function(inferred_images[b],
reference_images[b]).item()
predictions.append(
(loss, batch_index_offset + b, inferred_images[b]))
predictions.sort()
# Visualize or output the four worst and four best predictions
worst_filename = None if output_folder is None else output_folder + "/WorstPrediction.png"
best_filename = None if output_folder is None else output_folder + "/BestPrediction.png"
output_predictions(predictions[-5:-1], data_loader, "Worst predictions",
worst_filename)
output_predictions(predictions[0:4], data_loader, "Best predictions",
best_filename)
# Debug analyze code
if __name__ == '__main__':
from GlobalIlluminationFiltering import GlobalIlluminationFiltering
from Visualize import visualize_result
from ImageDataset import ImageDataset
from torch.utils.data import DataLoader
validation_set = ImageDataset(["Dataset/san-miguel/inputs"])
validation_data_loader = DataLoader(validation_set, batch_size=8)
model = GlobalIlluminationFiltering()
analyze_dataset(model, validation_data_loader)
class Config:
training_dir = "/home/hoang/comvis/datasets/at_t/orl_faces"
testing_dir = "/home/hoang/comvis/datasets/at_t/test"
train_batch_size = 32
train_num_epochs = 50
transforms = transforms.Compose(
[transforms.Resize((100, 100)),
transforms.ToTensor()])
folder_dataset = datasets.ImageFolder(root=Config.training_dir)
siamese_dataset = ImageDataset(imageFolderDataset=folder_dataset,
transform=transforms,
should_invert=False)
train_dataloader = DataLoader(dataset=siamese_dataset,
num_workers=2,
shuffle=True,
batch_size=Config.train_batch_size)
net = SiameseNetwork().cuda()
criterion = ContrastiveLoss()
optimizer = optim.Adam(net.parameters(), lr=0.005)
counter = []
loss_history = []
iteration_number = 0
