noise_reduce_cells.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ece7048/MHL-RGB-deep-learing/blob/main/noise_reduce_cells.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YIvx-8iMokKl"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "import subprocess\n",
        "from typing import Tuple, Dict\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "from torch import Tensor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "905sFSMlv7ug"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount(\"/content/gdrive\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EkKr2S91zMx1"
      },
      "outputs": [],
      "source": [
        "!pip install torch_geometric\n",
        "!pip install monai\n",
        "!pip install wandb\n",
        "!pip install quantus"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LPRchmoP2B2f"
      },
      "outputs": [],
      "source": [
        "loada3='resnet50'\n",
        "loada4='densenet121'\n",
        "loada='vgg16'\n",
        "loada2='vit'\n",
        "PATH= \"/content/gdrive/MyDrive/Colab/DATA_cell/trained_models/\"+loada2"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import monai\n",
        "from monai.networks.nets import DenseNet121, Classifier, EfficientNet, EfficientNetBN, ViT\n",
        "from torchvision import transforms, datasets, models\n",
        "\n",
        "class Net(nn.Module):\n",
        "\n",
        "    def __init__(self) -> None:\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = nn.Conv3d(1, 6, 5, 5)\n",
        "        self.pool = nn.MaxPool3d(2, 2)\n",
        "        self.conv2 = nn.Conv3d(6, 16, 2, 2)\n",
        "        self.fc1 = nn.Linear(1152 , 512)\n",
        "        self.fc2 = nn.Linear(512, 126)\n",
        "        self.fc3 = nn.Linear(126, 2)\n",
        "        #self.fc4 = nn.Linear(16, 2)\n",
        "\n",
        "    def forward(self, x: Tensor) -> Tensor:\n",
        "        x = self.pool(F.relu(self.conv1(x)))\n",
        "        x = self.pool(F.relu(self.conv2(x)))\n",
        "        x = x.view(x.shape[0],x.shape[1]*x.shape[2]*x.shape[3]*x.shape[4])\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = F.relu(self.fc3(x))\n",
        "       # x = F.relu(self.fc4(x))\n",
        "        return x\n",
        "\n",
        "class Net_2(nn.Module):\n",
        "\n",
        "    def __init__(self) -> None:\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = nn.Conv3d(1, 6, 5, 5)\n",
        "        self.pool = nn.MaxPool3d(2, 2)\n",
        "        self.conv2 = nn.Conv3d(6, 16, 2, 2)\n",
        "        self.fc1 = nn.Linear(1152 , 512)\n",
        "        self.fc2 = nn.Linear(512, 126)\n",
        "        self.fc3 = nn.Linear(126, 2)\n",
        "        #self.fc4 = nn.Linear(16, 2)\n",
        "        self.fcc1=nn.Linear(1024, 256)\n",
        "        self.fcc2=nn.Linear(256, 4)\n",
        "        self.fcc1=nn.Linear(4,2)\n",
        "\n",
        "\n",
        "    def forward(self, x: Tensor, y: Tensor) -> Tensor:\n",
        "        x = self.pool(F.relu(self.conv1(x)))\n",
        "        x = self.pool(F.relu(self.conv2(x)))\n",
        "        x = x.view(x.shape[0],x.shape[1]*x.shape[2]*x.shape[3]*x.shape[4])\n",
        "        x1 = F.relu(self.fc1(x))\n",
        "        x2 = F.relu(self.fc2(x1))\n",
        "        x3 = F.relu(self.fc3(x2))\n",
        "       # x = F.relu(self.fc4(x))\n",
        "        y = self.pool(F.relu(self.conv1(y)))\n",
        "        y = self.pool(F.relu(self.conv2(y)))\n",
        "        y = y.view(y.shape[0],x.shape[1]*y.shape[2]*y.shape[3]*y.shape[4])\n",
        "        y1 = F.relu(self.fc1(y))\n",
        "        y2 = F.relu(self.fc2(y1))\n",
        "        y3 = F.relu(self.fc3(y2))\n",
        "\n",
        "        c1=nn.cat((x1,y1))\n",
        "        yx1 = F.relu(self.fcc1(c1))\n",
        "        c2=nn.cat((x2,y2))\n",
        "        cc2=nn.cat((c2,yx1))\n",
        "        yx2=F.relu(self.fcc2(cc2))\n",
        "        c3=nn.cat((x3,y3))\n",
        "        cc3=nn.cat((c3,yx2))\n",
        "        yx3=F.relu(self.fcc3(cc3))\n",
        "\n",
        "        return yx3\n",
        "\n",
        "class CNNModel(nn.Module):\n",
        "    def __init__(self,num_classes):\n",
        "        super(CNNModel, self).__init__()\n",
        "        chn=1\n",
        "        self.conv_layer1 = self._conv_layer_set(chn, 32)\n",
        "        self.conv_layer2 = self._conv_layer_set(32, 64)\n",
        "        self.conv_layer3 = self._conv_layer_set(64, 128)\n",
        "        self.fc1 = nn.Linear(200704,1024)\n",
        "        self.fc2 = nn.Linear(1024, 512)\n",
        "        self.fc3 = nn.Linear(512, 256)\n",
        "        self.fc4 = nn.Linear(256, num_classes)\n",
        "        self.relu = nn.LeakyReLU()\n",
        "        self.batch=nn.BatchNorm1d(1024)\n",
        "        self.drop=nn.Dropout(p=0.3)\n",
        "\n",
        "    def _conv_layer_set(self, in_c, out_c):\n",
        "        conv_layer = nn.Sequential(\n",
        "        nn.Conv3d(in_c, out_c, kernel_size=(3, 3, 3), padding=0),\n",
        "        nn.LeakyReLU(),\n",
        "        nn.MaxPool3d((2, 2, 2)),\n",
        "        )\n",
        "        return conv_layer\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Set 1\n",
        "        x=x.float()\n",
        "        out = self.conv_layer1(x)\n",
        "        out = self.conv_layer2(out)\n",
        "        out = self.conv_layer3(out)\n",
        "        out = out.view(out.size(0), -1)\n",
        "        out = self.fc1(out)\n",
        "        out = self.relu(out)\n",
        "        if out.shape[0]>1:\n",
        "            out = self.batch(out)\n",
        "        out = self.drop(out)\n",
        "        out = self.fc2(out)\n",
        "        out = self.fc3(out)\n",
        "        out = self.fc4(out)\n",
        "        #print(out.shape)\n",
        "        return out\n",
        "\n",
        "\n",
        "class networks:\n",
        "    def __init__(self,name,device):\n",
        "        self.net=name\n",
        "        self.device=device\n",
        "\n",
        "    def build(self,c=2):\n",
        "        if self.net=='resnet50':\n",
        "            net=monai.networks.nets.resnet50(spatial_dims=2,n_input_channels=3,num_classes=c).to(self.device)\n",
        "        elif self.net=='vit':\n",
        "            net=monai.networks.nets.ViT(in_channels=3, img_size=(256,256), patch_size=(32,32), proj_type='conv', pos_embed_type='sincos', classification=True,\n",
        "num_classes=c, spatial_dims=2).to(self.device)\n",
        "        elif self.net=='densenet121':\n",
        "            net=monai.networks.nets.DenseNet121(spatial_dims=2, in_channels=3, out_channels=c).to(self.device)\n",
        "        elif self.net=='seresnext50':\n",
        "            net=monai.networks.nets.SEResNext50(spatial_dims=3, in_channels=1, out_channels=c).to(self.device)\n",
        "        elif self.net=='highresnet50':\n",
        "            net=monai.networks.nets.HighResNet(spatial_dims=3, in_channels=1, out_channels=c).to(self.device)\n",
        "        elif self.net=='vgg16':\n",
        "            net = models.vgg16(pretrained=False).to(self.device)\n",
        "            n_inputs = net.classifier[6].in_features\n",
        "            net.classifier[6] = nn.Sequential(nn.Linear(n_inputs, 256), nn.ReLU(), nn.Dropout(0.2),nn.Linear(256, c), nn.LogSoftmax(dim=1)).to(self.device)\n",
        "        else:\n",
        "            print('nodefine network!!')\n",
        "\n",
        "        return net"
      ],
      "metadata": {
        "id": "-NwO5RGf0PrB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torchdata\n",
        "%pip install vit-pytorch\n",
        "!pip install einops\n",
        "!pip install itk"
      ],
      "metadata": {
        "id": "E4CU9ovyp7hD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "DATA_ROOT1 = \"/content/gdrive/MyDrive/Colab/DATA_cell/A/\"\n",
        "DATA_ROOT2 = \"/content/gdrive/MyDrive/Colab/DATA_cell/B/\"\n",
        "\n",
        "\n",
        "import os\n",
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torchdata.datapipes.iter import IterableWrapper\n",
        "import nibabel as nib\n",
        "import numpy as np\n",
        "from scipy import ndimage\n",
        "import skimage.transform as skTrans\n",
        "from PIL import Image\n",
        "from monai.transforms import (\n",
        "    EnsureChannelFirstd,\n",
        "    AsDiscrete,\n",
        "    Compose,\n",
        "    CropForegroundd,\n",
        "    EnsureTyped,\n",
        "    FgBgToIndicesd,\n",
        "    LoadImaged,\n",
        "    Orientationd,\n",
        "    RandCropByPosNegLabeld,\n",
        "    ScaleIntensityRanged,\n",
        "    Spacingd,\n",
        "    RandFlip,\n",
        "    RandRotate,\n",
        "    RandZoom,\n",
        "    RandAffineGrid,\n",
        "    RandGaussianNoise,\n",
        "    RandShiftIntensity,\n",
        ")\n",
        "\n",
        "def normalize(volume):\n",
        "    min=np.min(volume)\n",
        "    max=np.max(volume)\n",
        "    volume[volume<min]=min\n",
        "    volume[volume>max]=max\n",
        "    volume=(volume-min)/(max-min)\n",
        "    volume=volume.astype(\"float32\")\n",
        "    return volume\n",
        "\n",
        "def resize_volume(img,d=84,w=128,h=128):\n",
        "    \"\"\"Resize across z-axis\"\"\"\n",
        "    # Set the desired depth\n",
        "    desired_depth = d\n",
        "    desired_width = w\n",
        "    desired_height = h\n",
        "    # Get current depth\n",
        "    current_depth = img.shape[2]\n",
        "    current_width = img.shape[0]\n",
        "    current_height = img.shape[1]\n",
        "\t\t# Compute depth factor\n",
        "    depth = current_depth / desired_depth\n",
        "    width = current_width / desired_width\n",
        "    height = current_height / desired_height\n",
        "    depth_factor = 1 / depth\n",
        "    width_factor = 1 / width\n",
        "    height_factor = 1 / height\n",
        "    #print(img.shape,desired_depth,desired_width,desired_height,width_factor,depth_factor,height_factor)\n",
        "    if len(img.shape)==4:\n",
        "        if img.shape[1]==img.shape[2]:\n",
        "            img=np.transpose(img, (1, 2, 0, 3))\n",
        "        img = ndimage.zoom(img, (width_factor, height_factor, depth_factor,1), order=1)\n",
        "    else:\n",
        "        if img.shape[1]==img.shape[2]:\n",
        "            img=np.transpose(img, (1, 2, 0))\n",
        "        img = ndimage.zoom(img, (width_factor, height_factor, depth_factor), order=1)\n",
        "    return img\n",
        "\n",
        "class GraphImageDataset(Dataset):\n",
        "    def __init__(self, data_path:str = \"\",\n",
        "                        transform=None,inp=0,sx=128,sy=128,sz=84):\n",
        "        self.data_path = data_path\n",
        "        self.transform = transform\n",
        "        self.inp=inp\n",
        "        self.sx,self.sy,self.sz=sx,sy,sz\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(os.listdir(self.data_path))-1\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        #image_collection.append(image)\n",
        "\n",
        "        # Load the image\n",
        "        sub=os.listdir(self.data_path)[index]\n",
        "        img_from_disk = Image.open(self.data_path+str(sub))\n",
        "        image = np.array(img_from_disk)\n",
        "        # Apply the transform to the image\n",
        "        #print(image.shape)\n",
        "        image=skTrans.resize(image,(256,256,3),order=1,preserve_range=True)\n",
        "        if self.transform:\n",
        "            image = self.transform((image))\n",
        "            #image = np.squeeze(image)\n",
        "        image=normalize(image)\n",
        "        #image=resize_volume(image,d=self.sz,w=self.sx,h=self.sy)\n",
        "        image=np.transpose(image,(2,0,1))\n",
        "        #image=skTrans.resize(image,(1,self.sx,self.sy,self.sz),order=1,preserve_range=True)\n",
        "        #print(image.shape)\n",
        "        image=torch.tensor(image)\n",
        "        y=torch.tensor(self.inp)\n",
        "        return image,y\n",
        "\n",
        "def load_data(batch_size=1, transform=None) -> Tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader, Dict]:\n",
        "    \"\"\"Load CIFAR-10 (training and test set).\"\"\"\n",
        "    # transform = Compose([RandGaussianNoise(prob=0.1),RandShiftIntensity(prob=0.2,offsets=(10,20)),RandFlip(prob=0.2),RandRotate(prob=0.2,range_x=[0.4,0.4])])\n",
        "    portion=0.8\n",
        "    trainset1 = GraphImageDataset(DATA_ROOT1, transform=transform, inp=0)\n",
        "    trainset2 = GraphImageDataset(DATA_ROOT2, transform=transform, inp=1)\n",
        "    print(len(trainset1),len(trainset2))\n",
        "    mid=int((len(trainset1)+len(trainset2))*portion)\n",
        "    whole=int((len(trainset1)+len(trainset2)))\n",
        "    print(mid,whole)\n",
        "    trainset = torch.utils.data.ConcatDataset([trainset1,trainset2])\n",
        "    gen1 = torch.Generator().manual_seed(whole)\n",
        "    whole_sampler = torch.utils.data.RandomSampler(trainset,num_samples=whole, generator=gen1)\n",
        "    indices=list(whole_sampler)\n",
        "    train_sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:mid])\n",
        "    test_sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[mid:])\n",
        "    print(len(trainset))\n",
        "    print(list(test_sampler))\n",
        "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,  sampler=train_sampler)\n",
        "    testloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,  sampler=test_sampler)\n",
        "    num_examples = {\"trainset\" : whole, \"testset\" : whole-mid}\n",
        "    return trainloader, testloader, num_examples, trainset, trainset\n",
        "\n"
      ],
      "metadata": {
        "id": "hqLSkwyTrKfl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "def data_build(bz=1):\n",
        "    print(\"Centralized PyTorch training\")\n",
        "    print(\"Load data\")\n",
        "    trainloader, testloader, _ ,_,_= load_data(batch_size=bz)\n",
        "    return trainloader, testloader\n"
      ],
      "metadata": {
        "id": "U2Hnh8ieshf7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256'\n",
        "%pip install \"monai--weekly[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil]\"\n"
      ],
      "metadata": {
        "id": "Otnp94E9YcX1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import PIL\n",
        "\n",
        "batch_n = 1\n",
        "\n",
        "train_loader, validation_loader= data_build(bz=batch_n)\n",
        "\n",
        "def imshow(img):\n",
        "    #img = img / 2 + 0.5     # unnormalize\n",
        "    npimg = img.numpy()\n",
        "    npimg=np.squeeze(npimg)\n",
        "    npimg=np.transpose(npimg, (1, 2, 0))\n",
        "    plt.imshow(npimg)\n",
        "    plt.imsave(PATH+'image.png',npimg)\n",
        "\n",
        "dataiter = iter(train_loader)\n",
        "images, labels = next(dataiter)\n",
        "\n",
        "# show images\n",
        "imshow(images)\n",
        "\n"
      ],
      "metadata": {
        "id": "1fIuqFQ0nxgP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import itk\n",
        "\n",
        "# Load the RGB image\n",
        "input_image = itk.imread(PATH+'image.png',itk.RGBPixel[itk.UC])\n",
        "# Define the pixel type and dimension for the image\n",
        "\n",
        "PixelType = itk.UC  # Unsigned Char for typical image data\n",
        "RGBPixelType = itk.RGBPixel[PixelType]\n",
        "Dimension = 2  # Assuming a 2D image\n",
        "\n",
        "# Define the image type\n",
        "ImageType = itk.Image[RGBPixelType, Dimension]\n",
        "\n",
        "# Cast the input image to the defined type\n",
        "cast_image = itk.CastImageFilter[input_image, ImageType].New(input_image)\n",
        "cast_image.Update()\n",
        "\n",
        "# Split the RGB image into its individual channels\n",
        "r_channel = itk.VectorIndexSelectionCastImageFilter.New(Input=input_image, Index=0)\n",
        "g_channel = itk.VectorIndexSelectionCastImageFilter.New(Input=input_image, Index=1)\n",
        "b_channel = itk.VectorIndexSelectionCastImageFilter.New(Input=input_image, Index=2)\n",
        "\n",
        "r_channel.Update()\n",
        "g_channel.Update()\n",
        "b_channel.Update()\n",
        "\n",
        "# Define the image type\n",
        "ImageType = itk.Image[PixelType, Dimension]\n",
        "\n",
        "# Define the Landweber Deconvolution filter type\n",
        "DeconvolutionFilterType = itk.LandweberDeconvolutionImageFilter[ImageType, ImageType]\n",
        "\n",
        "# Number of iterations for the deconvolution\n",
        "number_of_iterations = 50  # Adjust this as needed\n",
        "\n",
        "\n",
        "# Create a small Gaussian kernel\n",
        "kernel_size = [5, 5]  # Adjust the size as needed\n",
        "sigma = 1.0  # Standard deviation of the Gaussian\n",
        "\n",
        "# Create the image type for the PSF (2D image with float pixel type)\n",
        "ImageTypeF = itk.Image[itk.F, 2]\n",
        "ImageTypeUC = itk.Image[itk.UC, 2]\n",
        "\n",
        "# Create an empty image with the specified size and pixel type\n",
        "psf = ImageTypeF.New()\n",
        "psf.SetRegions(itk.Size[2](kernel_size))\n",
        "psf.Allocate()\n",
        "psf.FillBuffer(0.0)\n",
        "\n",
        "# Set a single pixel in the center to 1 (white spot)\n",
        "index = [int(k // 2) for k in kernel_size]\n",
        "psf.SetPixel(index, 1.0)\n",
        "\n",
        "# Apply Gaussian smoothing to create a Gaussian PSF\n",
        "gaussian_filter = itk.SmoothingRecursiveGaussianImageFilter[ImageTypeF, ImageTypeF].New()\n",
        "gaussian_filter.SetInput(psf)\n",
        "gaussian_filter.SetSigma(sigma)\n",
        "gaussian_filter.Update()\n",
        "\n",
        "kernel_image_float = gaussian_filter.GetOutput()\n",
        "\n",
        "# Cast the kernel image to unsigned char type\n",
        "caster = itk.CastImageFilter[ImageTypeF, ImageTypeUC].New()\n",
        "caster.SetInput(kernel_image_float)\n",
        "caster.Update()\n",
        "\n",
        "kernel_image = caster.GetOutput()\n",
        "\n",
        "# Set the KernelImage in the deconvolution filter\n",
        "deconvolve_r = DeconvolutionFilterType.New(Input=r_channel.GetOutput(), KernelImage=kernel_image, NumberOfIterations=number_of_iterations)\n",
        "deconvolve_g = DeconvolutionFilterType.New(Input=g_channel.GetOutput(), KernelImage=kernel_image, NumberOfIterations=number_of_iterations)\n",
        "deconvolve_b = DeconvolutionFilterType.New(Input=b_channel.GetOutput(), KernelImage=kernel_image, NumberOfIterations=number_of_iterations)\n",
        "\n",
        "deconvolve_r.Update()\n",
        "deconvolve_g.Update()\n",
        "deconvolve_b.Update()\n",
        "\n",
        "\n",
        "# Combine the deconvolved channels back into an RGB image\n",
        "rgb_image_deconvolved = itk.ComposeImageFilter.New(Input1=deconvolve_r.GetOutput(), Input2=deconvolve_g.GetOutput(), Input3=deconvolve_b.GetOutput())\n",
        "rgb_image_deconvolved.Update()\n",
        "\n",
        "# Save the result\n",
        "output_path = PATH+'image_filtered.png'\n",
        "itk.imwrite(rgb_image_deconvolved.GetOutput(), output_path)\n"
      ],
      "metadata": {
        "id": "3Y-M7q4CLr0P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def tuple_of_tensors_to_tensor(tuple_of_tensors):\n",
        "    return  torch.stack(list(tuple_of_tensors), dim=0)\n",
        "\n",
        "import warnings\n",
        "import os\n",
        "import wandb\n",
        "from monai.utils.enums import MetricReduction\n",
        "from monai.metrics import ConfusionMatrixMetric, ROCAUCMetric\n",
        "from monai.transforms import (\n",
        "    AsDiscrete,\n",
        "    Activations,\n",
        ")\n",
        "class AverageMeter(object):\n",
        "    def __init__(self):\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        self.val = 0\n",
        "        self.avg = 0\n",
        "        self.sum = 0\n",
        "        self.count = 0\n",
        "\n",
        "    def update(self, val, n=1):\n",
        "        self.val = val\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.avg = np.where(self.count > 0, self.sum / self.count, self.sum)\n",
        "\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "run_acc = AverageMeter()\n",
        "run_acc1 = AverageMeter()\n",
        "run_acc2 = AverageMeter()\n",
        "run_acc3 = AverageMeter()\n",
        "run_acc4 = AverageMeter()\n",
        "\n",
        "roc=ConfusionMatrixMetric(metric_name=\"accuracy\",compute_sample=True,reduction=MetricReduction.MEAN_BATCH, get_not_nans=False)\n",
        "f1=ConfusionMatrixMetric(metric_name=\"f1 score\",compute_sample=True,reduction=MetricReduction.MEAN_BATCH, get_not_nans=False)\n",
        "sens=ConfusionMatrixMetric(metric_name=\"sensitivity\",compute_sample=True,reduction=MetricReduction.MEAN_BATCH, get_not_nans=False)\n",
        "spec=ConfusionMatrixMetric(metric_name=\"specificity\",compute_sample=True,reduction=MetricReduction.MEAN_BATCH, get_not_nans=False)\n",
        "f1.reset()\n",
        "spec.reset()\n",
        "sens.reset()\n",
        "roc.reset()\n",
        "\n",
        "batch_n = 1\n",
        "\n",
        "train_loader, validation_loader= data_build(bz=batch_n)\n",
        "\n",
        "training=True  #turning it to TRue if you want training. If only testing False\n",
        "\n",
        "num_epochs = 100\n",
        "lr = 0.01\n",
        "name='vit'\n",
        "#net=CNNModel(num_classes=2).to(DEVICE)\n",
        "net2=networks(name,device)\n",
        "#net2=networks('vgg16',device)\n",
        "#net2=networks('resnet50',device)\n",
        "#net2=networks('densenet121',device)\n",
        "net=net2.build(c=2)\n",
        "\n",
        "loss_function = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(net.parameters(), lr=lr)\n",
        "post_label = AsDiscrete(argmax=True, to_onehot=2)\n",
        "\n",
        "\n",
        "#wandb.login(key=\"\")\n",
        "\n",
        "\n",
        "#run = wandb.init(\n",
        "#    # Set the project where this run will be logged\n",
        "#    project=\"ovarian-cancer-project\", #\"prova-project-ovarian\"\n",
        "#    # Track hyperparameters and run metadata\n",
        "#    config={\n",
        "#        \"learning_rate\": lr,\n",
        "#        \"epochs\": num_epochs,\n",
        "#        \"batch\" : batch_size\n",
        "#    })\n",
        "\n",
        "print_every = 1\n",
        "min=1000000000\n",
        "\n",
        "if os.path.exists(PATH+\".pth\"):\n",
        "    if torch.cuda.is_available():\n",
        "        checkpoint = torch.load(PATH+\".pth\")\n",
        "        net.load_state_dict(checkpoint['model_state_dict'])\n",
        "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
        "        min = checkpoint['loss']\n",
        "    else:\n",
        "        checkpoint = torch.load(PATH+\".pth\",map_location=torch.device('cpu'))\n",
        "        net.load_state_dict(checkpoint['model_state_dict'])\n",
        "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
        "        min = checkpoint['loss']\n",
        "    print('load last saved file: ',PATH, \" loss: \",min)\n",
        "\n",
        "if training==False:\n",
        "  num_epochs=0\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "\n",
        "    net.train()  # Set the model in training mode\n",
        "    total_loss_train = 0\n",
        "    correct_train = 0\n",
        "    total_train = 0\n",
        "\n",
        "    for i, data in enumerate(train_loader):\n",
        "        images=data[0].to(device=device, dtype=torch.float)\n",
        "        labels=data[1].to(device)\n",
        "\n",
        "        # Forward pass\n",
        "        output = net(images)\n",
        "        #output = torch.stack(list(output), dim=0)\n",
        "        if name=='vit':\n",
        "          output2=output[0]\n",
        "        else:\n",
        "          output2=output\n",
        "        output = torch.stack(list(output2), dim=0)\n",
        "        # Compute loss\n",
        "        loss = loss_function(output, labels)\n",
        "        total_loss_train += loss.item()\n",
        "\n",
        "        # Backward pass\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        # Calculate training accuracyfV\n",
        "        _, predicted = output.max(1)\n",
        "        correct_train += predicted.eq(labels).sum().item()\n",
        "        total_train += labels.size(0)\n",
        "\n",
        "\n",
        "    avg_loss_train = total_loss_train / len(train_loader)\n",
        "    accuracy_train = correct_train / total_train\n",
        "\n",
        "    net.eval()  # Set the model in evaluation mode\n",
        "    total_loss_val = 0\n",
        "    correct_val = 0\n",
        "    total_val = 0\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for i, data in enumerate(validation_loader):\n",
        "            images=data[0].to(device=device, dtype=torch.float)\n",
        "            label = data[1].to(device)\n",
        "            # Forward pass\n",
        "            output = net(images)\n",
        "            if name=='vit':\n",
        "                output2=output[0]\n",
        "            else:\n",
        "                output2=output\n",
        "            output1 = torch.stack(list(output2), dim=0)\n",
        "            labels=[post_label(out) for out in label]\n",
        "            output=[post_label(out1) for out1 in output1]\n",
        "            # Compute loss\n",
        "            loss = loss_function(output1, label)\n",
        "            total_loss_val += loss.item()\n",
        "            roc(y_pred=output, y=labels)\n",
        "            f1(y_pred=output, y=labels)\n",
        "            spec(y_pred=output, y=labels)\n",
        "            sens(y_pred=output, y=labels)\n",
        "            #print(output,labels)\n",
        "            # Calculate accuracy\n",
        "            _, predicted = output1.max(1)\n",
        "            correct_val += predicted.eq(label).sum().item()\n",
        "            total_val += label.size(0)\n",
        "            acc1 = roc.aggregate()\n",
        "            acc2_total = (f1.aggregate())\n",
        "            acc4_total = (spec.aggregate())\n",
        "            acc5_total = (sens.aggregate())\n",
        "            acc2=torch.nanmean(acc2_total[0])\n",
        "            acc4=torch.nanmean(acc4_total[0])\n",
        "            acc5=torch.nanmean(acc5_total[0])\n",
        "            acc6=torch.nanmean(acc1[0])\n",
        "            run_acc2.update(acc2.cpu().numpy())\n",
        "            run_acc1.update(acc6.cpu().numpy())\n",
        "            run_acc4.update(acc5.cpu().numpy())\n",
        "            run_acc3.update(acc4.cpu().numpy())\n",
        "            f1_v = run_acc2.avg\n",
        "            roc_v= run_acc1.avg\n",
        "            specif=run_acc3.avg\n",
        "            scens=run_acc4.avg\n",
        "\n",
        "    avg_loss_val = total_loss_val / len(validation_loader)\n",
        "    accuracy_val = correct_val / total_val\n",
        "    if avg_loss_val<min:\n",
        "            min=avg_loss_val\n",
        "            print('save model')\n",
        "            torch.save({'epoch': epoch,'model_state_dict': net.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'loss':avg_loss_val}, PATH+\".pth\")\n",
        "\n",
        "    print(f\"Epoch [{epoch+1}/{num_epochs}], Training Loss: {avg_loss_train:.4f}, Validation Loss: {avg_loss_val:.4f}, Train Accuracy: {accuracy_train:.4f}, Val Accuracy: {accuracy_val:.4f}\")\n",
        "    print(\"F1: \",f1_v,\"ROC: \",roc_v,\"Specif: \",specif, \"Scen: \",scens)\n",
        "\n",
        "    #wandb.log({\"loss_train\": avg_loss_train, \"loss_val\": avg_loss_val, \"acc_train\": accuracy_train, \"acc_val\": accuracy_val})\n",
        "\n",
        "\n",
        "#wandb.finish()\n"
      ],
      "metadata": {
        "id": "vUfNRdydw7X2"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "L4",
      "gpuClass": "premium",
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}