GSOC24_DeepLense_Test_5/plot_data.py at main · Commit2Cosmos/GSOC24_DeepLense_Test_5 · GitHub

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
import numpy as np
import matplotlib.pyplot as plt

import torch
from torchvision import transforms as transforms

# ds = "val"

# data = np.load(f"./data/{ds}/data_{ds}.npy")

# print(data[0].shape)

cols = 4

def plot_imgs_dataset(dataset, i, cols):
    fig, axes = plt.subplots(2, cols, sharex='all', sharey='all', figsize=(14,9))
    plt.axis('off')

    axes = axes.flatten()

    for j in range(0, len(axes)):

        im = dataset[i+j].squeeze()

        axes[j].imshow(im)


    plt.tight_layout()
    plt.show()


# for i in range(len(data)):
#     plot_imgs(data, i, cols)


from utils import get_loader_from_filenames
dataloader = get_loader_from_filenames('train', batch_size=2)

def plot_imgs_loader(tup):
    imgs, labels = tup

    _, axes = plt.subplots(1, imgs.shape[0], sharex='all', sharey='all', figsize=(14,9))
    plt.axis('off')

    axes = axes.flatten()

    for i in range(0, len(axes)):

        im = imgs[i].squeeze()

        axes[i].imshow(im)
        axes[i].set_title(f"Class {labels[i]}")


    plt.tight_layout()
    plt.show()


for d in dataloader:
    plot_imgs_loader(d)