<jupyter_start><jupyter_text>This notebook loads and displays training images with bounding box labels using PyTorch utility functions.Later on, you can enable some albumentations transformations and display the images with transformations applied.<jupyter_code>import os
import pprint

import matplotlib.pyplot as plt
from torchvision.utils import draw_bounding_boxes

from wheat.config import load_config
from wheat.data_module import WheatDataModule
from wheat import visualization as vis
%load_ext autoreload
%autoreload 2
%matplotlib inline
# change to the top level directory for this repository
os.chdir('..')
def display_image_grid(dataset, images_per_row=3, num_rows=3, first_image_index=0):
    for irow in range(num_rows):
        image_list = []
        for icol in range(images_per_row):
            image, labels = train_dataset[irow * images_per_row + icol + first_image_index]
            result = draw_bounding_boxes(
                vis.image_float_to_int_transform(image), labels['boxes'], colors='blue', width=5)
            image_list.append(result)
        vis.show(image_list)
# load the default configuration
config = load_config('wheat/config/config.ini')
pp = pprint.PrettyPrinter(indent=2)
# make sure data augmentation is off for first look at training images
for key in config['train']['transforms']:
    config['train']['transforms'][key] = 0  # set transform probability to 0
pp.pprint(config['train']['transforms'])
# initialize the dataset
wheat_data_module = WheatDataModule(config)
wheat_data_module.setup(stage='fit')
train_dataset = wheat_data_module.train_dataset
# display a few images
plt.rcParams['figure.figsize'] = [15, 10]
display_image_grid(train_dataset)
# turn on one or more transforms
config['train']['transforms']['color_jitter_prob'] = 1
# update the transforms associated with the dataset
wheat_data_module.config = config
train_dataset.transform = wheat_data_module.get_transforms()[0]
# display images with the transforms applied
display_image_grid(train_dataset)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08549096238397655, 'fraction_numerical': 0.006839276990718124, 'mean_word_length': 5.64935064935065, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '612474', 'license': 'MIT', 'n_tokens_mistral': 604, 'n_tokens_neox': 581, 'n_words': 172, 'path': 'notebooks/visualize_dataset.ipynb', 'repo_name': 'sheromon/wheat-detection', 'size': '2047'}