import FCN_Main
import torch
import numpy as np
from PIL import Image
import torch.nn.functional as F

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

model = torch.load('FCN')
model = model.to(device)
model = model.eval()

dir = "/home/rtx3090/storage/student1/cxb/nodecode/前期数据集/FCN_result/"

for i, data in enumerate(FCN_Main.test_loader):
    images, labels = data
    images, labels = images.to(device), labels.to(device)
    outputs = model(images)
    outputs = F.log_softmax(outputs, dim=1)
    pre_label = outputs.max(dim=1)[1].cpu().numpy()
    pre = Image.fromarray(np.uint8(pre_label.reshape(512,512)*255))
    pre.save(dir + str(i) + '.png')
    print('done')