def create_data_part(opt):
train_csv_path = os.path.join(opt.path_to_save_csv, 'train.csv')
val_csv_path = os.path.join(opt.path_to_save_csv, 'val.csv')
test_csv_path = os.path.join(opt.path_to_save_csv, 'test.csv')
train_ds = AVADataset(train_csv_path, opt.path_to_images, if_train = True)
val_ds = AVADataset(val_csv_path, opt.path_to_images, if_train = False)
test_ds = AVADataset(test_csv_path, opt.path_to_images, if_train=False)
train_loader = DataLoader(train_ds, batch_size=opt.batch_size, num_workers=opt.num_workers, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=opt.batch_size, num_workers=opt.num_workers, shuffle=False)
test_loader = DataLoader(test_ds, batch_size=opt.batch_size, num_workers=opt.num_workers, shuffle=False)
return train_loader, val_loader,test_loader
def get_data_loaders(csv_dir, image_dir, batch_size, num_workers):
train_ds = AVADataset(os.path.join(csv_dir, 'train.csv'), image_dir,
train_transform)
val_ds = AVADataset(os.path.join(csv_dir, 'val.csv'), image_dir,
val_transform)
test_ds = AVADataset(os.path.join(csv_dir, 'test.csv'), image_dir,
val_transform)
train_loader = DataLoader(train_ds,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
val_loader = DataLoader(val_ds,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
test_loader = DataLoader(test_ds,
batch_size=1,
num_workers=num_workers,
shuffle=False)
return train_loader, val_loader, test_loader
comment = sys.argv[1]
##############
# Dataloader #
##############
IMAGE_NET_MEAN = [0.485, 0.456, 0.406]
IMAGE_NET_STD = [0.229, 0.224, 0.225]
normalize = transforms.Normalize(mean=IMAGE_NET_MEAN, std=IMAGE_NET_STD)
train_dataset = AVADataset(
root_dir="/export/home/wangx/datasets/AVA/AVA_dataset/images/",
data_file="/export/home/wangx/datasets/AVA/AVA_dataset/AVA_filtered_train.txt",
transform=torchvision.transforms.Compose(
[
transforms.Resize((256, 256)),
transforms.RandomHorizontalFlip(),
transforms.RandomCrop((224, 224)),
transforms.ToTensor(),
normalize,
]
),
)
test_dataset = AVADataset(
root_dir="/export/home/wangx/datasets/AVA/AVA_dataset/images/",
data_file="/export/home/wangx/datasets/AVA/AVA_dataset/AVA_filtered_test.txt",
transform=torchvision.transforms.Compose(
[
transforms.Resize((224, 224)),
transforms.ToTensor(),
normalize,
]
from torchvision import transforms
from configs import DEVICE, NUM_EPOCHS, \
IMAGE_SIZE, \
DATASET_CONFIG, DATA_LOADER_CONFIG, \
MODEL_LOAD_PATH, EVAL_LOSS_PATH
from dataset import AVADataset, NIMADataset
from utils import plot_batch
from nima import NIMA, EMDLoss, vgg16_bn
# -- dataset --
# ava dataset
ava_datasets = {
phase:
AVADataset(image_root=DATASET_CONFIG["image_root"],
label_filepath=DATASET_CONFIG[phase + "_label_filepath"])
for phase in TEST_PHASE_LIST
}
# transforms
transforms = {
TEST:
transforms.Compose(transforms=[
transforms.Resize(IMAGE_SIZE),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(1, 1, 1))
])
}
nima_datasets = {
phase: NIMADataset(dataset=ava_datasets[phase],
ground_truth = torch.from_numpy(ground_truth)
return image, ground_truth
#%%
if __name__ == "__main__":
from torchvision import transforms
from torch.utils.data import DataLoader
from dataset import AVADataset
from configs import AVA_LABEL_CLEANED_FILEPATH, AVA_IMAGE_ROOT
from utils import plot_image_ratings, plot_batch
ava_dataset = AVADataset(image_root=AVA_IMAGE_ROOT,
label_filepath=AVA_LABEL_CLEANED_FILEPATH)
print(ava_dataset.labels.describe())
transforms = transforms.Compose([
transforms.Resize(256), transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
# transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(1, 1, 1))
])
nima_dataset = NIMADataset(
dataset=ava_dataset, transforms=transforms,
is_target_tensor=True
)
index = np.random.choice(range(len(nima_dataset)))
#%%
import numpy as np
import pickle
from configs import EVAL_PRED_PATH, TEST, DATASET_CONFIG
from dataset import AVADataset
from utils import ave_rating, plot_prediction
with open(EVAL_PRED_PATH, "rb") as fr:
predictions = pickle.load(fr)[TEST]
# ava dataset
ava_dataset = AVADataset(image_root=DATASET_CONFIG["image_root"],
label_filepath=DATASET_CONFIG[TEST +
"_label_filepath"])
#%%
# instance
index = np.random.randint(low=0, high=len(ava_dataset))
image = ava_dataset[index]["image"]
ground_truth = ava_dataset[index]["ground_truth"]
image_id = predictions.loc[index, AVADataset._label_key_image_id]
prediction = predictions.loc[index, AVADataset._label_key_ratings]
plot_prediction(image=image, ground_truth=ground_truth, prediction=prediction)
print("image id: {}".format(predictions.loc[index, "image_id"]))
print("predicted average score: {}".format(ave_rating(prediction)))
print("ground truth score: {}".format(ave_rating(ground_truth)))
#%%
epochs = 50
log_step = 50
learning_rate = 3e-6
comment = ""
##############
# Dataloader #
##############
IMAGE_NET_MEAN = [0.485, 0.456, 0.406]
IMAGE_NET_STD = [0.229, 0.224, 0.225]
normalize = transforms.Normalize(mean=IMAGE_NET_MEAN, std=IMAGE_NET_STD)
train_dataset = AVADataset(transform=torchvision.transforms.Compose([
transforms.Resize((256, 256)),
transforms.RandomHorizontalFlip(),
transforms.RandomCrop((224, 224)),
transforms.ToTensor(),
normalize,
]))
test_dataset = AVADataset(transform=torchvision.transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
normalize,
]), )
train_loader = DataLoader(train_dataset,
batch_size,
shuffle=True,
num_workers=4)
test_loader = DataLoader(test_dataset, 128, shuffle=False, num_workers=4)
##########