The following code
from torch.utils.data import DataLoader
from defect_segmentation.data_loading.DatasetSingleImage import dataset_single_image_default
import numpy as np
if __name__ == "__main__":
import matplotlib.pyplot as plt
def main():
dataset = dataset_single_image_default()
batch_size = 16
shuffle = True
num_workers = 0
loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
for i_batch, sample_batched in enumerate(loader):
fig, axs = plt.subplots(int(np.sqrt(batch_size)), batch_size // int(np.sqrt(batch_size)))
fig.suptitle(f"i_batch = {i_batch}")
for i_sample, ax in zip(range(sample_batched.shape[0]), axs.flat):
ax.set_title(f"Sample #{i_sample}")
ax.axis("off")
ax.imshow(sample_batched[i_sample, :, :])
plt.pause(0.001)
main()
works and outputs figures like so
This is fine.
Problem is by the 10th figure, it becomes very slow to populate the figures. I don't know what can be causing this.
For completeness, here is the code for creating the dataset (striding over a single image):
from torch.utils.data import Dataset
from Utils.ConfigProvider import ConfigProvider
import cv2
import os
from overrides import overrides # pip install overrides
class DatasetSingleImage(Dataset):
def __init__(self, image_path: str, sample_shape: tuple, strides: tuple):
self._path = image_path
assert os.path.isfile(self._path)
self._im = cv2.imread(self._path)
self._shape = self._im.shape
self._rows, self._cols = self._shape[0], self._shape[1]
self._sample_shape = sample_shape
self._sample_rows, self._sample_cols = self._sample_shape[0], self._sample_shape[1]
self._strides = strides
self._stride_rows, self._stride_cols = self._strides[0], self._strides[1]
# self._rows_start_range = range(0, self._rows, self._stride_rows)
# self._cols_start_range = range(0, self._cols, self._stride_cols)
self._rows_tuples_range = \
[(c, min(c + self._sample_rows, self._rows)) for c in range(0, self._rows - self._sample_rows, self._stride_rows)]
self._cols_tuples_range = \
[(r, min(r + self._sample_cols, self._cols)) for r in range(0, self._cols - self._sample_cols, self._stride_cols)]
self._n_strides_rows = len(self._rows_tuples_range)
self._n_strides_cols = len(self._cols_tuples_range)
self._total_strides = self._n_strides_rows * self._n_strides_cols
def __len__(self):
return self._total_strides
@overrides # pip install overrides
def __getitem__(self, ind):
row_ind = ind // self._n_strides_cols
col_ind = ind % self._n_strides_cols
sample_x = self._rows_tuples_range[row_ind]
sample_y = self._cols_tuples_range[col_ind]
sample = self._im[sample_x[0]:sample_x[1], sample_y[0]:sample_y[1]]
assert sample.shape[:2] == self._sample_shape
return sample
def dataset_single_image_default():
path = ConfigProvider.config().data.defective_inspected_path1
sample_shape = (50, 50)
strides = (25, 25)
dataset = DatasetSingleImage(path, sample_shape, strides)
return dataset
What is making the plotting slow, and how to fix it?