카테고리 없음

jpg make_grid.py

notou10 2022. 8. 3. 03:06

한장씩 저장 -> torchvision.save_image가 더 편함

from torchvision import transforms
import torchvision
from torchvision.utils import save_image

for index, (data, label) in tqdm(enumerate(dataloader), total=len(dataloader)): 
    #cuda 올리기
    data = data.to('cuda')
    
    ################
    #case 1. 좌상단  
    ori = data.detach()
    parted = data.detach()[:, :, :128, :128]
    ori[:, :, :128, 128:] = parted
    #################
    
    for i in range(bs):
        idx = bs*index + i 
        filename = f'{str(idx).zfill(5)}.png'
        save_image(ori[i], os.path.join(out_dir, filename))

 

plt로 학습 중 1장 보기

plt.close()
plt.imshow(ori[0].permute(1, 2, 0).detach().cpu().numpy())
plt.show()
plt.savefig('./d.png')

 

 

 

make_grid

from torchvision import utils
import torch
from natsort import natsorted
import os
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt

path = "../../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/fsmr_dog_adained"
out_dir = os.path.join(path[:-len(path.split("/")[-1])], "grid")

#import pdb; pdb.set_trace()
if not os.path.isdir(out_dir):
    os.mkdir(out_dir)
    

a = natsorted(os.listdir(path))
transform = transforms.ToTensor()
output_tensor_array = torch.zeros(16, 3, 256, 256)

for index, img_path in enumerate(natsorted(os.listdir(path))[:16]):
    img = Image.open(os.path.join(path, img_path))
    img = transform(img)
    img.to('cuda')
    output_tensor_array[index] = img
    

grid = utils.make_grid(output_tensor_array, nrow = 4)


plt.imshow(grid.permute(1, 2, 0))
plt.axis('off')
plt.savefig(os.path.join(out_dir,path.split("/")[-1]), dpi = 600)

 

import numpy as np
import os
import torch
from PIL import Image
from torchvision import transforms
from torchvision import utils
import matplotlib.pyplot as plt

name = "stylegan"
grid_img_num = 16

path = "../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/fsmr_dog_adained_all/"
list_a = np.load('fsmr_npy/vit_fsmr.npy')
list_b = np.load('fsmr_npy/resnet_fsmr.npy')
#list_b = np.load('fsmr_npy/vit_stylegan.npy')


sorted_dog = np.sort(list_a[:,0]-list_b[:,0])
sort_index_dog = np.argsort(list_a[:,0]-list_b[:,0])

dog_out_dir = f"fsmr_npy/{name}_-predicted-dog" #fsmr 은 dog 이라 예상
cat_out_dir = f"fsmr_npy/{name}_-predicted-cat" #fsmr 은 cat 이라 예상

transform = transforms.ToTensor()

if not os.path.isdir(dog_out_dir):
    os.mkdir(dog_out_dir)
    os.mkdir(cat_out_dir)
    

output_tensor_array = torch.zeros(grid_img_num, 3, 256, 256)
plt.close()
for index, data in enumerate(range(grid_img_num)):
    img = Image.open(os.path.join(path, f"{sort_index_dog[index]}.png"))
    img = transform(img)
    img.to('cuda')
    output_tensor_array[index] = img

grid = utils.make_grid(output_tensor_array, nrow = 4)

            

plt.imshow(grid.permute(1, 2, 0))
for x in range(4):
    for y in range(4):
        plt.text(200+ 256*x, 225+ 256*y, f"img{x+4*y}", ha="center", va="center", size=10)
        
plt.axis('off')
plt.savefig(os.path.join(dog_out_dir,"dog"), dpi = 600)

plt.close()
plt.bar(range(grid_img_num), sorted_dog[:16])
plt.title("vit predicted this image as cat , resnet did dog")
plt.savefig(os.path.join(dog_out_dir,"dog_bar"), dpi = 600)