see also:
from pathlib import Path
# Setup path to a data folder
data_path = Path("data/")
image_path = data_path / "main folder name for images"
# Setup train and testing paths
train_dir = image_path / "train"
test_dir = image_path / "test"
import os
def walk_through_dir(dir_path):
"""Walks through dir_path returning its contents."""
for dirpath, dirnames, filenames in os.walk(dir_path):
print(f"There are {len(dirnames)} directories and {len(filenames)} images in '{dirpath}'.")
walk_through_dir(image_path)
import random
from PIL import Image
# Set seed
# random.seed(42)
# 1. Get all image paths
image_path_list = list(image_path.glob("*/*/*.jpg"))
# 2. Pick a random image path
random_image_path = random.choice(image_path_list)
# 3. Get image class from path name (the image class is the name of the directory where the image is stored)
image_class = random_image_path.parent.stem
# 4. Open image
img = Image.open(random_image_path)
# 5. Print metadata
print(f"Random image path: {random_image_path}")
print(f"Image class: {image_class}")
print(f"Image height: {img.height}")
print(f"Image width: {img.width}")
img
import numpy as np
import matplotlib.pyplot as plt
# Turn the image into an array
img_as_array = np.asarray(img)
# Plot the image with matplotlib
plt.figure(figsize=(10, 7))
plt.imshow(img_as_array)
plt.title(f"Image class: {image_class} | Image shape: {img_as_array.shape} -> [height, width, color_channels] (HWC)")
plt.axis(False);
#optionally view the image array: img_as_array
import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms # Write a transform for image data_transform = transforms.Compose([ # Resize our images to 64x64 transforms.Resize(size=(64, 64)), # Flip the images randomly on the horizontal transforms.RandomHorizontalFlip(p=0.5), # Turn the image into a torch.Tensor transforms.ToTensor() ])
def plot_transformed_images(image_paths: list, transform, n=3, seed=None):
"""
Selects random images from a path of images and loads/transforms
them then plots the original vs the transformed version.
"""
if seed:
random.seed(seed)
random_image_paths = random.sample(image_paths, k=n)
for image_path in random_image_paths:
with Image.open(image_path) as f:
fig, ax = plt.subplots(nrows=1, ncols=2)
ax[0].imshow(f)
ax[0].set_title(f"Original\nSize: {f.size}")
ax[0].axis(False)
# Transform and plot target image
transformed_image = transform(f).permute(1, 2, 0) # note we will need to change shape for matplotlib (C, H, W) -> (H, W, C)
ax[1].imshow(transformed_image)
ax[1].set_title(f"Transformed\nShape: {transformed_image.shape}")
ax[1].axis("off")
fig.suptitle(f"Class: {image_path.parent.stem}", fontsize=16)
plot_transformed_images(image_paths=image_path_list,
transform=data_transform,
n=3,
seed=None)
# Use ImageFolder to create dataset(s)
from torchvision import datasets
train_data = datasets.ImageFolder(root=train_dir,
transform=data_transform, # a transform for the data
target_transform=None) # a transform for the label/target
test_data = datasets.ImageFolder(root=test_dir,
transform=data_transform)
train_data, test_data
import os
import pathlib
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from typing import Tuple, Dict, List
# Instance of torchvision.datasets.ImageFolder()
train_data.classes, train_data.class_to_idx
# Setup path for target directory
target_directory = train_dir
print(f"Target dir: {target_directory}")
# Get the class names from the target directory
class_names_found = sorted([entry.name for entry in list(os.scandir(target_directory))])
class_names_found
list(os.scandir(target_directory))
def find_classes(directory: str) -> Tuple[List[str], Dict[str, int]]:
"""Finds the class folder names in a target directory."""
# 1. Get the class names by scanning the target directory
classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
# 2. Raise an error if class names could not be found
if not classes:
raise FileNotFoundError(f"Couldn't find any classes in {directory}... please check file structure.")
# 3. Create a dictionary of index labels (computers prefer numbers rather than strings as labels)
class_to_idx = {class_name: i for i, class_name in enumerate(classes)}
return classes, class_to_idx
find_classes(target_directory)
# 0. Write a custom dataset class
from torch.utils.data import Dataset
# 1. Subclass torch.utils.data.Dataset
class ImageFolderCustom(Dataset):
# 2. Initialize our custom dataset
def __init__(self,
targ_dir: str,
transform=None):
# 3. Create class attributes
# Get all of the image paths
self.paths = list(pathlib.Path(targ_dir).glob("*/*.jpg"))
# Setup transform
self.transform = transform
# Create classes and class_to_idx attributes
self.classes, self.class_to_idx = find_classes(targ_dir)
# 4. Create a function to load images
def load_image(self, index: int) -> Image.Image:
"Opens an image via a path and returns it."
image_path = self.paths[index]
return Image.open(image_path)
# 5. Overwrite __len__()
def __len__(self) -> int:
"Returns the total number of samples."
return len(self.paths)
# 6. Overwrite __getitem__() method to return a particular sample
def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
"Returns one sample of data, data and label (X, y)."
img = self.load_image(index)
class_name = self.paths[index].parent.name # expects path in format: data_folder/class_name/image.jpg
class_idx = self.class_to_idx[class_name]
# Transform if necessary
if self.transform:
return self.transform(img), class_idx # return data, label (X, y)
else:
return img, class_idx # return untransformed image and label
# Create a transform
from torchvision import transforms
train_transforms = transforms.Compose([
transforms.Resize(size=(64, 64)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor()
])
test_transforms = transforms.Compose([
transforms.Resize(size=(64, 64)),
transforms.ToTensor()
])
# Test out ImageFolderCustom
train_data_custom = ImageFolderCustom(targ_dir=train_dir,
transform=train_transforms)
test_data_custom = ImageFolderCustom(targ_dir=test_dir,
transform=test_transforms)
# Get class names as list
class_names = train_data.classes
class_names
class_dict = train_data.class_to_idx
class_dict
# Check the lengths of our dataset
len(train_data), len(test_data)
train_data.samples[0]
# Index on the train_data Dataset to get a single image and label
img, label = train_data[0][0], train_data[0][1]
print(f"Image tensor:\n {img}")
print(f"Image shape: {img.shape}")
print(f"Image datatype: {img.dtype}")
print(f"Image label: {label}")
print(f"Label datatype: {type(label)}")
#display resized image:
# Rearrange the order dimensions
img_permute = img.permute(1, 2, 0)
# Print out different shapes
print(f"Original shape: {img.shape} -> [color_channels, height, width]")
print(f"Image permute: {img_permute.shape} -> [height, width, color_channels]")
# Plot the image
plt.figure(figsize=(10, 7))
plt.imshow(img_permute)
plt.axis("off")
plt.title(class_names[label], fontsize=14)
import os
os.cpu_count()
# Turn train and test datasets into DataLoader's
from torch.utils.data import DataLoader
BATCH_SIZE=1
train_dataloader = DataLoader(dataset=train_data,
batch_size=BATCH_SIZE,
num_workers=1,
shuffle=True)
test_dataloader = DataLoader(dataset=test_data,
batch_size=BATCH_SIZE,
num_workers=1,
shuffle=False)
train_dataloader, test_dataloader
len(train_dataloader), len(test_dataloader)
img, label = next(iter(train_dataloader))
# Batch size will now be 1, you can change the batch size if you like
print(f"Image shape: {img.shape} -> [batch_size, color_channels, height, width]")
print(f"Label shape: {label.shape}")
# 1. Create a function to take in a dataset
def display_random_images(dataset: torch.utils.data.Dataset,
classes: List[str] = None,
n: int = 10,
display_shape: bool = True,
seed: int = None):
# 2. Adjust display if n is too high
if n > 10:
n = 10
display_shape = False
print(f"For display, purposes, n shouldn't be larger than 10, setting to 10 and removing shape display.")
# 3. Set the seed
if seed:
random.seed(seed)
# 4. Get random sample indexes
random_samples_idx = random.sample(range(len(dataset)), k=n)
# 5. Setup plot
plt.figure(figsize=(16, 8))
# 6. Loop through random indexes and plot them with matplotlib
for i, targ_sample in enumerate(random_samples_idx):
targ_image, targ_label = dataset[targ_sample][0], dataset[targ_sample][1]
# 7. Adjust tensor dimensions for plotting
targ_image_adjust = targ_image.permute(1, 2, 0) # [color_channels, height, width] -> [height, width, color_channels]
# Plot adjusted samples
plt.subplot(1, n, i+1)
plt.imshow(targ_image_adjust)
plt.axis("off")
if classes:
title = f"Class: {classes[targ_label]}"
if display_shape:
title = title + f"\nshape: {targ_image_adjust.shape}"
plt.title(title)
# Display random images from the ImageFolder created Dataset
display_random_images(train_data,
n=5,
classes=class_names,
seed=None)
# Display random images from the ImageFolderCustom Dataset
display_random_images(train_data_custom,
n=5,
classes=class_names,
seed=None)
# Let's look at trivailaugment - https://pytorch.org/vision/stable/auto_examples/plot_transforms.html#trivialaugmentwide
from torchvision import transforms
train_transform = transforms.Compose([
transforms.Resize(size=(224, 224)),
transforms.TrivialAugmentWide(num_magnitude_bins=31),
transforms.ToTensor()
])
test_transform = transforms.Compose([
transforms.Resize(size=(224, 224)),
transforms.ToTensor()
])