def same_class(
query_dataset_name: str,
queried_dataset_name: str,
query_embeddings_name: str,
queried_embeddings_name: str,
k: int,
distance: str,
n_gpu: int,
processes: int,
):
"""Cross-modal image recall of same class elements."""
click.echo(
"Calculating cross modal class recall@{} for the {} and {} embeddings".
format(k, query_embeddings_name, queried_embeddings_name))
query_dataset, queried_dataset = (
get_dataset(query_dataset_name),
get_dataset(queried_dataset_name),
)
query_embeddings = load_embeddings(query_embeddings_name)
queried_embeddings = load_embeddings(queried_embeddings_name)
average_class_recall_parallel(
query_dataset,
queried_dataset,
query_embeddings,
queried_embeddings,
k,
distance,
n_gpu,
processes,
)
def retrieve_top_k(
model,
query_image_filename,
query_dataset_name,
queried_dataset_name,
queried_embeddings_name,
k=16,
distance="cosine",
n_gpu=0,
):
"""Query the embeddings for a dataset with the given image. The image is embedded with the given model. Pairwise
distances to the query image are computed for each embedding in the dataset, so the embeddings created by `model`
must have the same length as the ones in the embedding directory. The `k` most similar images are displayed.
Args:
model: name of the model to be used for embedding the DataSet.
query_image_filename: the complete file path and name to the image to be used as query. The images in the
dataset that are most similar to this image will be displayed.
query_dataset_name: name of the registered dataset to which the query image belongs.
queried_dataset_name: name of the registered dataset to be queried.
queried_embeddings_name: the name of the pickle file where the embeddings will be saved.
k: the number of most similar images that wil be displayed. (Default value = 16)
distance: which distance function to be used for nearest neighbor computation. Either 'cosine' or 'pairwise' (Default value = "cosine")
n_gpu: number of available GPUs. If zero, the CPU will be used. (Default value = 0)
Returns:
"""
device = get_device(n_gpu)
# Load data
query_dataset = get_dataset(query_dataset_name)
queried_dataset = get_dataset(queried_dataset_name)
# Load embeddings from pickle directory
queried_embeddings = load_embeddings(queried_embeddings_name).to(device)
# Get the query image and create the embedding for it
image, image_class = query_dataset[query_image_filename]
# Send elements to the specified device
image, model = [var.to(device) for var in [image, model]]
query_embedding = model(image.unsqueeze(
0)) # unsqueeze to add the missing dimension expected by the model
# Compute the distance to the query embedding for all images in the Dataset
queried_embeddings, query_embedding = [
var.to(device) for var in [queried_embeddings, query_embedding]
]
top_distances, top_indices = get_top_k(query_embedding, queried_embeddings,
k, distance)
plot_image_retrieval(image, image_class, query_dataset, queried_dataset,
top_distances, top_indices)
def create_embeddings(model, dataset_name, embeddings_name, batch_size,
workers, n_gpu):
"""Creates embedding vectors for each element in the given DataSet, and saves a single Pytorch tensor of shape
`len(dataset), embeddings_size` with all the embeddings in a file with the given name.
Args:
model: name of the model to be used for embedding the DataSet.
dataset_name: name of the registered dataset which will be embedded.
embeddings_name: the name of the pickle directory where the embeddings will be saved.
batch_size: size of batches for the embedding process.
workers: number of data loader workers.
n_gpu: number of available GPUs. If zero, the CPU will be used.
Returns:
"""
device = get_device(n_gpu)
model = model.eval().to(device)
dataset = get_dataset(dataset_name)
data_loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=workers)
embedding_list = [
model(data[0].to(device)
) # model output sent to 'cpu' to prevent gpu memory overflow
for data in tqdm(
data_loader, total=len(data_loader), desc="Embedding data")
]
torch.save(
torch.cat(embedding_list).to("cpu"),
open(get_path("embeddings", "{}.pt".format(embeddings_name)), "wb"))
def same_class(
dataset_name: str,
embeddings_name: str,
test_split: float,
k: int,
distance: str,
n_gpu: int,
) -> None:
"""Image recall of same class elements."""
click.echo("Calculating class recall@{} for {} embeddings".format(
k, embeddings_name))
dataset = get_dataset(dataset_name)
embeddings = load_embeddings(embeddings_name)
(
query_embeddings,
queried_embeddings,
query_indices,
queried_indices,
) = random_simple_split(embeddings, test_split)
query_dataset, queried_dataset = (
Subset(dataset, query_indices),
Subset(dataset, queried_indices),
)
average_class_recall(
query_dataset,
queried_dataset,
query_embeddings,
queried_embeddings,
k,
distance,
n_gpu,
)
def initial_model(self):
""" """
dataset = get_dataset(self.dataset_name)
image_dimension = dataset[0][0].shape[1]
return HOGGCN(
self.dataset.classes_dataframe,
image_dimension,
self.in_channels,
self.cell_size,
self.bins,
self.signed_gradients,
self.processes,
)
def plot_embedding_tsne(dataset_name, embeddings_name, load_projection=False):
"""Plot a 2D projection of embeddings in the specified embedding directory using plotly.
Args:
dataset_name: name of the registered dataset which will be embedded.
embeddings_name: the name of the directory where the batch pickles will be saved.
load_projection: load projections from pickles. (Default value = False)
Returns:
"""
from vscvs.embeddings import load_embeddings # import here to avoid circular import
dataset = get_dataset(dataset_name)
embeddings = load_embeddings(embeddings_name).to("cpu")
projection_pickle_dir = get_path("embeddings", embeddings_name)
if load_projection:
click.echo("Loading existing 2D projection from pickle.")
projection = pickle.load(
open(os.path.join(projection_pickle_dir, "tsne.pickle"), "rb"))
dataset_class_names = pickle.load(
open(
os.path.join(projection_pickle_dir, "tsne_class_names.pickle"),
"rb"))
else:
click.echo("Creating 2D projection of the embeddings using TSNE")
projection = TSNE(n_components=2).fit_transform(embeddings)
dataset_class_names = [
dataset.classes[tup[1]]
for tup in tqdm(dataset, desc="Retrieving image class names")
]
pickle.dump(
projection,
open(os.path.join(projection_pickle_dir, "tsne.pickle"), "wb"))
pickle.dump(
dataset_class_names,
open(
os.path.join(projection_pickle_dir, "tsne_class_names.pickle"),
"wb"))
trace = go.Scattergl( # plot the resulting projection using plotly
x=projection[:, 0],
y=projection[:, 1],
text=dataset_class_names,
mode="markers",
marker=dict(size=16,
color=np.random.randn(len(projection)),
colorscale="Viridis"),
)
data = [trace]
plot(data)
def display_sample_batch(dataset_name, batch_size, workers):
"""Output a random batch form the specified dataset.
Args:
dataset_name: name of the registered dataset which will be embedded.
batch_size: size of batches for the embedding process.
workers: number of data loader workers.
Returns:
"""
dataset = get_dataset(dataset_name) # load data
data_loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=workers) # create the data_loader
plot_image_batch(next(iter(data_loader))) # plot the batch
def create_metadata_tsv(dataset_name):
"""Creates a metadata TSV file that can be used with any embeddings of the given dataset to be displayed in the
Tensorboard embedding projector.
Args:
dataset_name: name of the registered DatasetFolder which will be embedded.
Returns:
"""
dataset = get_dataset(dataset_name)
with open(get_path("embeddings", "{}.tsv".format(dataset_name)), "w") as f:
writer = csv.DictWriter(f, ["class_idx", "class", "path"],
delimiter="\t")
writer.writeheader()
writer.writerows([{
"class_idx": idx,
"class": dataset.classes[idx],
"path": path
} for path, idx in tqdm(dataset.samples, desc="Creating tsv")])
def sprite_tensor(dataset_name, size=(64, 64)):
"""Create an image tensor that can be used to create a Tensorboard embedding projector sprite image for the dataset
corresponding to `dataset_name`. This image tensor must have a shape of `(N, C, H, W)` where `N` is the length of
the dataset, `C` is the number of color channels, and `H` and `W` are the height and width dimensions for the
sprite components. Dataset elements must be reduced in size in order to fit in memory when using Tensorboard.
Args:
dataset_name: the name of the dataset from which the sprite tensor will be created.
size: the dimension of the resulting sprite image components. If `size` is a sequence like `(h, w)`, the
output size will be matched to this. If size is an int, the smaller edge of the image will be matched to this
number. i.e, if `height > width`, then image will be rescaled to `(size * height / width, size)`. (Default value = (64)
64):
Returns:
the sprite image tensor, with the dimensions specified above.
"""
from vscvs.datasets import get_dataset
dataset = get_dataset(dataset_name, normalize=False, size=size)
data_loader = DataLoader(dataset, batch_size=len(dataset), shuffle=False)
return next(iter(data_loader))[0]
def initial_model(self):
""" """
dataset = get_dataset(self.dataset_name)
return GCNClassification(len(dataset.classes), 11)
def sketchy_suite(
dataset_name: str,
photo_embeddings_name: str,
sketch_embeddings_name: str,
test_split: float,
top_k: List[int],
distance: str,
n_gpu: int,
) -> None:
"""Same mode recall for photos and sketches, as well as cross-modal recall using
both modes as queries.
"""
photo_dataset = get_dataset(f"{dataset_name}-photos")
sketch_dataset = get_dataset(f"{dataset_name}-sketches")
photo_embeddings = load_embeddings(photo_embeddings_name)
sketch_embeddings = load_embeddings(sketch_embeddings_name)
for k in top_k:
click.echo(f"\nCalculating class recall@{k}\n\nphoto to photo:")
(
query_embeddings,
queried_embeddings,
query_indices,
queried_indices,
) = random_simple_split(photo_embeddings, test_split)
query_dataset, queried_dataset = (
Subset(photo_dataset, query_indices),
Subset(photo_dataset, queried_indices),
)
average_class_recall( # photo to photo average class recall
query_dataset,
queried_dataset,
query_embeddings,
queried_embeddings,
k,
distance,
n_gpu,
)
click.echo(f"\nsketch to sketch:")
(
query_embeddings,
queried_embeddings,
query_indices,
queried_indices,
) = random_simple_split(sketch_embeddings, test_split)
query_dataset, queried_dataset = (
Subset(sketch_dataset, query_indices),
Subset(sketch_dataset, queried_indices),
)
average_class_recall( # sketch to sketch average class recall
query_dataset,
queried_dataset,
query_embeddings,
queried_embeddings,
k,
distance,
n_gpu,
)
click.echo(f"\nphoto to sketch:")
average_class_recall( # photo to sketch average class recall
photo_dataset,
sketch_dataset,
photo_embeddings,
sketch_embeddings,
k,
distance,
n_gpu,
)
click.echo(f"\nsketch to photo:")
average_class_recall( # sketch to photo average class recall
sketch_dataset,
photo_dataset,
sketch_embeddings,
photo_embeddings,
k,
distance,
n_gpu,
)
def __init__(self,
*args,
batch_size=0,
dataset_name=None,
drop_last=False,
epochs=1,
n_gpu=0,
parameter_dict=None,
resume_date=None,
resume_checkpoint=None,
tags=None,
train_validation_split=0.8,
workers=6,
**kwargs):
"""
:param args: mixin arguments.
:type: Tuple
:param batch_size: batch size during training.
:type: int
:param dataset_name: the name of the Dataset to be used for training.
:type: str
:param drop_last: whether to drop the last batch if it is not the same size as `batch_size`.
:type: bool
:param epochs: the number of epochs used for training.
:type: int
:param n_gpu: number of GPUs available. Use 0 for CPU mode.
:type: int
:param parameter_dict: dictionary with important training parameters for logging.
:type: Dict
:param resume_date: date of the trainer state to be resumed. Dates must have this format: `%y-%m-%dT%H-%M`.
:type: str
:param resume_checkpoint: name of the model checkpoint to be loaded.
:type: str
:param tags: optional tags for model checkpoint and tensorboard logs organization.
:type: List[int]
:param train_validation_split: proportion of the training set that will be used for actual
training. The remaining data will be used as the validation set.
:type: float
:param workers: number of workers for data_loader.
:type: int
:param kwargs: mixin keyword arguments.
:type: Dict
"""
date = datetime.now()
self.batch_size = batch_size
self.checkpoint_directory = get_checkpoint_path(self.trainer_id,
*tags,
date=date)
self.dataset = get_dataset(dataset_name)
self.dataset_name = dataset_name
self.device = get_device(n_gpu)
self.epochs = epochs
self.log_directory = get_log_directory(self.trainer_id,
*tags,
date=date)
self.model = self.initial_model
self.parameter_dict = parameter_dict
self.resume_date = datetime.strptime(
resume_date,
CHECKPOINT_NAME_FORMAT) if resume_date else resume_date
self.resume_checkpoint = resume_checkpoint
self.start_epoch = 1
self.tags = tags
self._load_checkpoint()
self.epoch = self.start_epoch
self.step = 0
self.train_loader, self.validation_loader = self._create_data_loaders(
train_validation_split, batch_size, workers, drop_last)
self.trainer_engine = self._create_trainer_engine()
self.evaluator_engine = self._create_evaluator_engine()
self.timer = self._create_timer()
self.progressbar = self._progressbar
self.writer = SummaryWriter(self.log_directory)
self._add_event_handlers()
self._add_model_checkpoint_savers()
super().__init__(*args, **kwargs)
