# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + id="iTPm3f0pkxkI" colab_type="code" outputId="c18940a4-b612-4998-91f4-5e8704e4c164" colab={"base_uri": "https://localhost:8080/", "height": 171}
on_google_colab = False
model_availaible = True
if on_google_colab:
  from google.colab import drive
  drive.mount('/content/drive')
  # !ls '/content/drive/My Drive'

# + id="kfsYl9NTJ-wp" colab_type="code" outputId="ca41f97e-452e-4194-fce0-78de632640f4" colab={"base_uri": "https://localhost:8080/", "height": 101}
# !pip install pydicom

# + id="NkEufWFIETr9" colab_type="code" colab={}
import torch
import numpy as np
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler

# + id="j9gwC6Z6mnk9" colab_type="code" outputId="25e37368-37c9-4b41-aaa0-253c41c12f69" colab={"base_uri": "https://localhost:8080/", "height": 50}
from load_LIDC_data import LIDC_IDRI

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_dataset = LIDC_IDRI(dataset_location = '/content/drive/My Drive/')
# test_dataset = LIDC_IDRI(dataset_location = '/content/drive/My Drive/', training=False)
dataset_size = len(train_dataset)
indices = list(range(dataset_size))
split = int(np.floor(0.1 * dataset_size))
np.random.shuffle(indices)
train_indices, test_indices = indices[split:], indices[:split]
train_sampler = SubsetRandomSampler(train_indices)
test_sampler = SubsetRandomSampler(test_indices)
train_loader = DataLoader(train_dataset, batch_size=8, sampler=train_sampler)
test_loader = DataLoader(train_dataset, batch_size=1, sampler=test_sampler)
print("Number of training/test patches:", (len(train_indices),len(test_indices)))

# + id="Y6zUvWppzDQJ" colab_type="code" colab={}
from importlib import reload
import model
# reload(model)

net = model.ProbU_Net(in_channels=1, out_channels=1, latent_dim=6)
net.to(device)
optimizer = torch.optim.Adam(net.parameters())

# + id="69HbCPnynYMd" colab_type="code" outputId="b8d04baa-f704-4e5f-ddde-c885b415186e" colab={"base_uri": "https://localhost:8080/", "height": 1000}
loss_history = []
import random
for epoch in range(15):
    train_loss = 0.0
    for step, (patch, mask, _) in enumerate(train_loader):
        # print(mask.shape) 
        mask = mask[:,random.randint(0,3),:,:]
        patch = patch.to(device)
        mask = mask.to(device)
        mask = torch.unsqueeze(mask,1)
        net.forward(patch, mask, train=True)
        loss = net.loss(mask, beta=10.0)
        x = loss.item()
        train_loss += x
        if (step%100==0):
          print('(%d, %5d) item-loss: %.9f'%(epoch + 1, step + 1, x))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print('[%d] epoch-loss: %.9f'%(epoch + 1, train_loss))
    loss_history.append(train_loss)
print('Finished Training')

# + id="DodwfkUMFwLr" colab_type="code" outputId="db36aae5-f659-4f25-ba94-7d7e2c940af0" colab={"base_uri": "https://localhost:8080/", "height": 286}
import matplotlib.pyplot as plt
plt.plot(loss_history)

# + id="zt41Q_dnqXPw" colab_type="code" outputId="51a23767-fceb-4bf7-9882-7038d160fe07" colab={"base_uri": "https://localhost:8080/", "height": 179}
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torchvision

def sigmoid(array):
  return 1/(1+np.exp(-array))

ns = 8
for step, (patch, masks, _) in enumerate(test_loader):
    patch = patch.to(device)
    recon = [sigmoid(x.cpu().detach().numpy()) for x in net.reconstruct(patch, num_samples=ns)]
    masks = masks.cpu().detach().numpy()
    patch = patch.cpu().detach().numpy()
    # print(masks.shape)
    f = plt.figure()
    f.add_subplot(1,5,1)
    plt.imshow(patch[0][0],cmap='gray')
    for i in range(4):
      f.add_subplot(1,5,i+2)
      plt.imshow(masks[0,i,:,:],cmap='gray')
    
    g = plt.figure()
    for i in range(ns):
      g.add_subplot(1,ns,i+1)
      plt.imshow(recon[i][0][0],cmap='gray')

    break


# + id="D9GNeTwAhiZo" colab_type="code" colab={}
PATH = 'punet.pth'
torch.save(net.state_dict(), PATH)

# + id="Eajp05kr5oyG" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 1000} outputId="0eba713a-fbcf-44f0-f8da-350c89b3fe50"
import model
PATH = 'punet.pth'
net = model.ProbU_Net(in_channels=1, out_channels=1, latent_dim=6)
net.load_state_dict(torch.load(PATH))
net.cuda()

# + id="IVziypPb6v90" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 470} outputId="50f3fba5-f316-4058-e2f5-20dbb962364a"
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torchvision

def sigmoid(array):
  return 1/(1+np.exp(-array))

ns = 6
for step, (patch, masks, _) in enumerate(test_loader):
    patch = patch.to(device)
    # net.reconstruct(patch, num_samples=ns)
    recon = [sigmoid(x.cpu().detach().numpy()) for x in net.reconstruct(patch, num_samples=ns)]
    masks = masks.cpu().detach().numpy()
    patch = patch.cpu().detach().numpy()
    # print(masks.shape)
    f = plt.figure(figsize=(20, 16))
    f.add_subplot(1,5,1)
    plt.imshow(patch[0][0],cmap='gray')
    for i in range(4):
      f.add_subplot(1,5,i+2)
      plt.imshow(masks[0,i,:,:],cmap='gray')
    
    recon = np.round(recon)
    g = plt.figure(figsize=(30, 24))
    for i in range(ns):
      g.add_subplot(1,ns,i+1)
      plt.imshow(recon[i][0][0],cmap='gray')

    break

StarcoderdataJupyterScriptsDedupFiltered