- Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathvis.py
More file actions
Latest commit
35 lines (31 loc) · 1.23 KB
/
Copy pathvis.py
File metadata and controls
35 lines (31 loc) · 1.23 KB
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
importos
importrandom
importmatplotlib.pyplotasplt
importnumpyasnp
importtorch
fromproj2_code.dl_utilsimportpredict_labels
fromproj2_code.image_loaderimportImageLoader
defvisualize(model: torch.nn.Module,
split: str,
data_transforms,
data_base_path: str='../data') ->None:
loader=ImageLoader(data_base_path, split=split, transform=data_transforms)
class_labels=loader.class_dict
class_labels= {ele.lower(): class_labels[ele] foreleinclass_labels}
labels= {class_labels[ele]: eleforeleinclass_labels}
paths_and_labels=loader.load_imagepaths_with_labels(class_labels)
selected=random.choices(paths_and_labels, k=4)
fig, axs=plt.subplots(2, 2)
foriinrange(4):
img=loader.load_img_from_path(selected[i][0])
withtorch.no_grad():
outputs=model(data_transforms(img).unsqueeze(
0).to(next(model.parameters()).device))
predicted=predict_labels(outputs).item()
axs[i//2, i%2].imshow(img, cmap='gray')
axs[i//2, i%2].set_title('Predicted:{}|Correct:{}'.format(
labels[predicted], labels[selected[i][1]]))
axs[i//2, i%2].axis('off')
fig.tight_layout()
plt.subplots_adjust(wspace=0.5)
plt.show()