Spaces:
Runtime error
Runtime error
File size: 4,582 Bytes
b78dbf0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import numpy as np
import os
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
from PIL import Image
from typing import List, Tuple
def make_dataset(path: str) -> Tuple[List[str], List[str]]:
"""
Creates a dataset of paired images from a directory.
The dataset should be partitioned into two sets: one contains images that
will have the low pass filter applied, and the other contains images that
will have the high pass filter applied.
Args
- path: string specifying the directory containing images
Returns
- images_a: list of strings specifying the paths to the images in set A,
in lexicographically-sorted order
- images_b: list of strings specifying the paths to the images in set B,
in lexicographically-sorted order
"""
images_a = []
images_b = []
all_images = os.listdir(path)
for image in all_images:
if image[1] == 'a':
images_a.append(os.path.join(path, image))
else:
images_b.append(os.path.join(path, image))
images_a, images_b = np.sort(images_a), np.sort(images_b)
return images_a, images_b
def get_cutoff_frequencies(path: str) -> List[int]:
"""
Gets the cutoff frequencies corresponding to each pair of images.
The cutoff frequencies are the values you discovered from experimenting in
part 1.
Args
- path: string specifying the path to the .txt file with cutoff frequency
values
Returns
- cutoff_frequencies: numpy array of ints. The array should have the same
length as the number of image pairs in the dataset
"""
cutoff_frequencies = []
freq_file = open(path, 'r')
freqs = freq_file.readlines()
for freq in freqs:
freq = int(freq.strip())
cutoff_frequencies.append(freq)
cutoff_frequencies = np.array(cutoff_frequencies)
return cutoff_frequencies
class HybridImageDataset(data.Dataset):
"""Hybrid images dataset."""
def __init__(self, image_dir: str, cf_file: str) -> None:
"""
HybridImageDataset class constructor.
You must replace self.transform with the appropriate transform from
torchvision.transforms that converts a PIL image to a torch Tensor. You can
specify additional transforms (e.g. image resizing) if you want to, but
it's not necessary for the images we provide you since each pair has the
same dimensions.
Args:
- image_dir: string specifying the directory containing images
- cf_file: string specifying the path to the .txt file with cutoff
frequency values
"""
images_a, images_b = make_dataset(image_dir)
cutoff_frequencies = get_cutoff_frequencies(cf_file)
self.transform = transforms.Compose([transforms.ToTensor()])
self.images_a = images_a
self.images_b = images_b
self.cutoff_frequencies = cutoff_frequencies
def __len__(self) -> int:
"""Returns number of pairs of images in dataset."""
return len(self.images_a)
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, int]:
"""
Returns the pair of images and corresponding cutoff frequency value at
index `idx`.
Since self.images_a and self.images_b contain paths to the images, you
should read the images here and normalize the pixels to be between 0 and 1.
Make sure you transpose the dimensions so that image_a and image_b are of
shape (c, m, n) instead of the typical (m, n, c), and convert them to
torch Tensors.
Args
- idx: int specifying the index at which data should be retrieved
Returns
- image_a: Tensor of shape (c, m, n)
- image_b: Tensor of shape (c, m, n)
- cutoff_frequency: int specifying the cutoff frequency corresponding to
(image_a, image_b) pair
HINTS:
- You should use the PIL library to read images
- You will use self.transform to convert the PIL image to a torch Tensor
"""
image_a_dir = self.images_a[idx]
image_b_dir = self.images_b[idx]
cutoff_frequency = self.cutoff_frequencies[idx]
image_a = Image.open(image_a_dir)
image_b = Image.open(image_b_dir)
pixels_a = np.array(image_a) / 255.0
pixels_b = np.array(image_b) / 255.0
image_a = self.transform(pixels_a).float()
image_b = self.transform(pixels_b).float()
return image_a, image_b, cutoff_frequency
|