Projeto-DACO-V4-UNET-1HOT-Wei-CellPose.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Check GPU ","metadata":{"id":"TAvWegoHHh4i"}},{"cell_type":"code","source":"!nvidia-smi","metadata":{"id":"AXWDGRGzHjzu","outputId":"f9d33237-081c-4769-8fce-849cab4d2f68","execution":{"iopub.status.busy":"2022-02-02T11:00:12.4001Z","iopub.execute_input":"2022-02-02T11:00:12.400442Z","iopub.status.idle":"2022-02-02T11:00:13.108537Z","shell.execute_reply.started":"2022-02-02T11:00:12.400373Z","shell.execute_reply":"2022-02-02T11:00:13.107722Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Cellpose images as labels","metadata":{"id":"Fzj8vgZKBHcG"}},{"cell_type":"code","source":"\nimport numpy as np \nimport matplotlib.pyplot as plt\nimport skimage\nfrom skimage import io\nfrom PIL import Image\nfrom sklearn.metrics import jaccard_score\nimport os\n\ncellpose=os.listdir('/content/drive/MyDrive/DACO/CellPose')\ncellpose=[x for x in cellpose if x.endswith('.npy')]\nsave_path='/content/drive/MyDrive/DACO/CellPose_Flow_Png'\nfor x in cellpose:\n  seg=np.load(os.path.join('/content/drive/MyDrive/DACO/CellPose',x),allow_pickle=True).item()\n  array=np.array(seg['flows'])[0][0]\n  array=array.clip(min=0)\n  array=array.clip(max=255) \n  seg_img=Image.fromarray(np.uint8(array))\n  gray = seg_img.convert('L')  #conversion to gray scale    \n  gray.save(os.path.join(save_path,x[:8]+'_flow.png'))\n  \n  del seg\n\n\n","metadata":{"id":"olz5rZGT5B0F"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset Class","metadata":{"id":"HW7j9idQPJdQ"}},{"cell_type":"code","source":"import numpy as np\nimport csv\nfrom PIL import Image\nimport os\ncsvfile=open(r\"../input/sartorius-cell-instance-segmentation/train.csv\")\n\ncsvreader=csv.reader(csvfile)\nids=[]\nclass_=[]\ncsv=[]\nfor row in csvreader:\n  csv.append(row);\n  ids.append(row[0])\n  class_.append(row[4])\n\nclass_=np.array(class_).reshape(-1,1)\nids=np.array(ids).reshape(-1,1)\n\nclasses=np.concatenate([ids,class_],axis=1)\n\n\ntrain=os.listdir('../input/sartorius-cell-instance-segmentation/train')\n\ntrain=[x[:-4] for x in train]\ncort_ids=[classes[x][0] for x in range(classes.shape[0]) if classes[x][1]=='cort']\nastro_ids=[classes[x][0] for x in range(classes.shape[0]) if classes[x][1]=='astro']\nshsy5y_ids=[classes[x][0] for x in range(classes.shape[0]) if classes[x][1]=='shsy5y']\ncort_ids_f=[]\nastro_ids_f=[]\nshsy5y_ids_f=[]\n\n[cort_ids_f.append(n) for n in cort_ids if n not in cort_ids_f] \n[astro_ids_f.append(n) for n in astro_ids if n not in astro_ids_f] \n[shsy5y_ids_f.append(n) for n in shsy5y_ids if n not in shsy5y_ids_f] \n\ncort=np.array(['cort' for x in range(len(cort_ids_f))]).reshape(-1,1)\nastro=np.array(['astro' for x in range(len(astro_ids_f))]).reshape(-1,1)\nshsy5y=np.array(['shsy5y' for x in range(len(shsy5y_ids_f))]).reshape(-1,1)\n\ncort_ids_f=np.array(cort_ids_f).reshape(-1,1)\nastro_ids_f=np.array(astro_ids_f).reshape(-1,1)\nshsy5y_ids_f=np.array(shsy5y_ids_f).reshape(-1,1)\n\n\ncort=np.concatenate([cort_ids_f,cort],axis=1)\nastro=np.concatenate([astro_ids_f,astro],axis=1)\nshsy5y=np.concatenate([shsy5y_ids_f,shsy5y],axis=1)\n\nids_classes=np.concatenate([cort,astro,shsy5y])\n\n\n","metadata":{"id":"mbm-EtNMChd0","execution":{"iopub.status.busy":"2022-02-03T16:18:26.790095Z","iopub.execute_input":"2022-02-03T16:18:26.791154Z","iopub.status.idle":"2022-02-03T16:18:28.299738Z","shell.execute_reply.started":"2022-02-03T16:18:26.791031Z","shell.execute_reply":"2022-02-03T16:18:28.298863Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport numpy as np\nimport torch\nimport skimage\nfrom skimage import io\nfrom PIL import Image\n\nroot_path = '../input/sartorius-cell-instance-segmentation'\n\n#Class which allow iterate through images from the dataset\nclass SartoriusDataset(torch.utils.data.Dataset):\n    def __init__(self, root,mask_root, transforms,id_classes):\n        self.root = root\n        self.mask_root=mask_root\n        self.transforms = transforms\n        self.id_classes=id_classes\n        #Load all training images and masks, sorting them to ensure that they are aligned\n        self.imgs = list(sorted(os.listdir(os.path.join(root, \"train\"))))\n        self.imgs=self.imgs[:-1]\n        self.masks = list(sorted(os.listdir(os.path.join(mask_root, \"CellPose_Flow_Png\"))))\n\n    def __getitem__(self, idx):\n        #Load images and masks\n        \n        img_path = os.path.join(self.root, \"train\", self.imgs[idx])\n        mask_path = os.path.join(self.mask_root, \"CellPose_Flow_Png\", self.masks[idx])\n\n        #Open image and mask\n        img = io.imread(img_path)\n        #img = skimage.color.gray2rgb(img)\n        mask = io.imread(mask_path)\n        \n        zeros=np.expand_dims(np.zeros((520,704)),axis=-1)\n        ones=np.expand_dims(np.ones((520,704)),axis=-1)\n       \n        if ids_classes[idx][1]=='cort':\n            img = np.expand_dims(img, axis=-1)\n            img=np.concatenate([img,ones,zeros,zeros],axis=-1)\n        if ids_classes[idx][1]=='astro':\n            img = np.expand_dims(img, axis=-1)\n            img=np.concatenate([img,zeros,ones,zeros],axis=-1)\n        if ids_classes[idx][1]=='shsy5y':\n            img = np.expand_dims(img, axis=-1)\n            img=np.concatenate([img,zeros,zeros,ones],axis=-1)\n        \n        if self.transforms is not None:\n            img =self.transforms(img)\n            mask = self.transforms(mask)\n        \n        \n        #return img, target\n        return img, mask\n\n    def __len__(self):\n        return len(self.imgs)\n\n#Class which allow iterate through images from the dataset\nclass SartoriusTestDataset(torch.utils.data.Dataset):\n    def __init__(self, root, transforms):\n        self.root = root\n        self.transforms = transforms\n\n        #Load all training images and masks, sorting them to ensure that they are aligned\n        self.imgs = list(sorted(os.listdir(os.path.join(root, \"test\"))))\n\n    def __getitem__(self, idx):\n        #Load images and masks\n        img_path = os.path.join(self.root, \"test\", self.imgs[idx])\n\n        #Open image and mask\n        img = io.imread(img_path)\n        #img = skimage.color.gray2rgb(img)\n        \n        ones=np.expand_dims(np.ones((520,704)),axis=-1)\n        zeros=np.expand_dims(np.zeros((520,704)),axis=-1)\n        if ids_classes[idx][1]=='cort':\n            img = np.expand_dims(img, axis=-1)\n            img=np.concatenate([img,ones,zeros,zeros],axis=-1)\n        if ids_classes[idx][1]=='astro':\n            img = np.expand_dims(img, axis=-1)\n            img=np.concatenate([img,zeros,ones,zeros],axis=-1)\n        if ids_classes[idx][1]=='shsy5y':\n            img = np.expand_dims(img, axis=-1)\n            img=np.concatenate([img,zeros,zeros,ones],axis=-1)\n        if self.transforms is not None:\n            img =self.transforms(img)\n\n        #return img, target\n        return img\n\n    def __len__(self):\n        return len(self.imgs)","metadata":{"id":"yQk4PnDbRS4x","execution":{"iopub.status.busy":"2022-02-03T16:18:38.605377Z","iopub.execute_input":"2022-02-03T16:18:38.605855Z","iopub.status.idle":"2022-02-03T16:18:40.807692Z","shell.execute_reply.started":"2022-02-03T16:18:38.605816Z","shell.execute_reply":"2022-02-03T16:18:40.806738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# U-Net Model","metadata":{"id":"k24D1d9pvP7W"}},{"cell_type":"code","source":"import torch\nfrom torch import nn\nfrom torchvision import models\n\n\nclass UNET(nn.Module):\n    def __init__(self,in_channels,init_features):\n        super().__init__()         \n        \n        features=init_features\n        \n        # Construct the encoder blocks\n        self.enc1 = double_conv(in_channels, features)       \n        self.enc2 =  double_conv(features, features * 2)      \n        self.enc3 = double_conv(features * 2, features * 4)       \n        self.enc4 =double_conv(features * 4, features * 8)       \n        self.bottleneck = b_conv(features * 8, features * 16)   \n        self.pool=nn.MaxPool2d(kernel_size=2, stride=2)\n\n        # Construct the decoder blocks\n        self.dec1=dec(features*8)\n        self.upconv1=upscale(features*16,features*8)\n        self.dec2=dec(features*4)\n        self.upconv2=upscale(features*8,features*4)\n        self.dec3=dec(features*2)\n        self.upconv3=upscale(features*4,features*2)\n        self.dec4=dec(features)\n        self.upconv4=upscale(features*2,features)\n        self.conv=nn.Conv2d(features,1, kernel_size=1) #out\n    \n    def forward(self,x): #x input image\n        # encoder\n        enc1 = self.enc1(x)\n        \n        enc2 = self.enc2(self.pool(enc1))\n        \n        enc3 = self.enc3(self.pool(enc2))\n        \n        enc4 = self.enc4(self.pool(enc3))\n        \n        bottleneck = self.bottleneck(self.pool(enc4))\n        \n        dec1=self.upconv1(bottleneck)\n        dec1=copy_crop(dec1,enc4)\n        dec1=self.dec1(dec1)\n        \n        dec2=self.upconv2(dec1)\n        dec2=copy_crop(dec2,enc3)\n        dec2=self.dec2(dec2)\n        \n        dec3=self.upconv3(dec2)\n        dec3=copy_crop(dec3,enc2)\n        dec3=self.dec3(dec3)\n        \n        dec4=self.upconv4(dec3)\n        dec4=copy_crop(dec4,enc1)\n        dec4=self.dec4(dec4)\n        \n        out=self.conv(dec4) \n        \n        return out\n    def to_class(output):\n        output=[1 if (x>0.5) else 0 for x in output]\n  \n        return output\ndef dec(feat):\n    return nn.Sequential(\n            nn.Conv2d(feat*2,feat,kernel_size=3,padding=1,bias=False),\n            nn.BatchNorm2d(feat),\n            nn.ReLU(),\n            nn.Conv2d(feat,feat,kernel_size=3,padding=1,bias=False),\n            nn.BatchNorm2d(feat),\n            nn.ReLU()\n            )\n    \ndef upscale (in_f,out_f): \n    return nn.ConvTranspose2d(in_f, out_f,kernel_size=2,stride=2)\n    \ndef copy_crop (dec,enc): \n    return torch.cat((dec,enc),dim=1)\n\ndef double_conv(in_c, out_c):\n    conv_=nn.Sequential(\n        nn.Conv2d(in_c, out_c, kernel_size=3,padding=1,bias=False),\n        nn.BatchNorm2d(out_c),\n        nn.ReLU(),\n        nn.Conv2d(out_c, out_c, kernel_size=3,padding=1,bias=False),\n        nn.BatchNorm2d(out_c),\n        nn.ReLU()\n        )\n    return conv_\n\n\n\ndef b_conv(in_c, out_c):\n    conv=nn.Sequential(\n        nn.Conv2d(in_c, out_c, kernel_size=3,padding=1),\n        nn.ReLU()\n        )\n    return conv","metadata":{"id":"FN4yf3JPvSmA","execution":{"iopub.status.busy":"2022-02-03T16:18:48.044561Z","iopub.execute_input":"2022-02-03T16:18:48.045119Z","iopub.status.idle":"2022-02-03T16:18:48.249324Z","shell.execute_reply.started":"2022-02-03T16:18:48.045081Z","shell.execute_reply":"2022-02-03T16:18:48.248615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training Loop","metadata":{"id":"di4yyBTrvLEq"}},{"cell_type":"code","source":"import numpy as np\nimport torch\nfrom PIL import Image\nfrom torch import nn\nimport skimage\nfrom skimage import io\nfrom skimage import transform\nimport torchvision.transforms as transforms\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import SubsetRandomSampler\nfrom tqdm import tqdm\nimport copy\n\npil = transforms.ToPILImage()\n\ndef IoU(outputs, labels):\n  intersection = np.logical_and(labels, outputs)\n  union = np.logical_or(labels, outputs)\n  iou_score = np.sum(intersection) / np.sum(union)\n\n  return iou_score\n\n#Path to dataset\nroot_path = '../input/sartorius-cell-instance-segmentation'\nsave_path='./'\nmask_root='../input/cellpose-flow/content/drive/MyDrive/DACO'\n#Create transforms and compose\ncomposed_transform = transforms.Compose([transforms.ToTensor(),transforms.Resize((512,512)) ])\n\n#Dataset settings\ndataset = SartoriusDataset(root_path,mask_root, composed_transform,ids_classes)\n\nidx = [*range(len(dataset))]\ntraining_idx, validation_idx = train_test_split(idx, test_size = 0.2)\n\ntraining_sampler = SubsetRandomSampler(training_idx)\nvalidation_sampler = SubsetRandomSampler(validation_idx)\n\nbatch_size = 2\ntraining_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=training_sampler, num_workers=0)\nvalidation_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=validation_sampler, num_workers=0)\n\n#Define model\nmodel = UNET(4,32)\n\ncriterion =nn.MSELoss()\n\n#Define one optimizer and scheduler for learning rate\noptimizer = torch.optim.SGD(model.parameters(), 0.01 , momentum=0.01)\nscheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.5)\n#Check if GPU is available for training\ntrain_on_gpu = torch.cuda.is_available()\n\nif train_on_gpu:\n  print(\"Cuda is available! Training on GPU\")\n  model.cuda()\nelse:\n  print(\"Cuda isn't available! Training on CPU\")\n\n#Number of epochs\nn_epoch = 50;\n\ntraining_IoU = []\ntraining_loss = []\n\nvalidation_IoU = []\nvalidation_loss = []\nlrs = []\nmin_valid_loss = 1e9\n\nfor epoch in range(n_epoch):\n  print(f'[Epoch: {epoch+1}]')\n\n  #Training loop\n  model.train()\n  print('Training model...')\n\n  t_loss = 0\n  t_IoU = 0\n\n  for i, (img,mask) in enumerate(tqdm(training_dataloader)):\n    \n    if train_on_gpu:\n      img, mask = img.cuda(), mask.cuda()\n    \n    #Reset the optimizer gradient\n    optimizer.zero_grad()\n    img=img.float()\n    #Feedforward\n    output = model(img)\n    \n    #Calculate the loss\n    loss = criterion(torch.sigmoid(output), mask)\n\n    #Backpropagation\n    loss.backward()\n\n    #Update the model\n    optimizer.step()\n    \n    #Save loss\n    t_loss += loss.item()\n\n    #Save IoU\n    for j in range(len(output)):\n        t_IoU += IoU(np.round(torch.sigmoid(output)[j][0].detach().cpu().numpy()+0.3),np.round(mask[j][0].detach().cpu().numpy()+0.3))\n        \n    #Visualize Image:\n    '''\n    out=pil(torch.round(torch.sigmoid(output[0])))\n    display(out)'''\n\n  #Validation loop\n  model.eval()\n  print('Validating model...')\n\n  v_loss = 0\n  v_IoU = 0\n\n  for i, (img,mask) in enumerate(tqdm(validation_dataloader)):\n    \n    if train_on_gpu:\n      img, mask = img.cuda(), mask.cuda()\n\n    #Feedforward\n    img=img.float()\n    output = model(img)\n    \n    #Calculate the loss\n    loss = criterion(torch.sigmoid(output), mask)\n\n    #Save loss\n    v_loss += loss.item()\n\n    #Save IoU\n    for j in range(len(output)):\n      v_IoU += IoU(np.round(torch.sigmoid(output)[j][0].detach().cpu().numpy()+0.3),np.round(mask[j][0].detach().cpu().numpy()+0.3))\n    \n  # learning rate update every 10 epochs:\n  if (epoch%10==0): \n    scheduler.step()\n    lrs.append(optimizer.param_groups[0][\"lr\"])\n    print(f'Learning Rate Updated to {lrs[-1]}')\n  #Average and save losses and IoU metric\n  t_loss = t_loss/len(training_dataloader.sampler)\n  v_loss = v_loss/len(validation_dataloader.sampler)\n  t_IoU = t_IoU/len(training_dataloader.sampler)\n  v_IoU = v_IoU/len(validation_dataloader.sampler)\n\n  training_loss.append(t_loss)\n  validation_loss.append(v_loss)\n  training_IoU.append(t_IoU)\n  validation_IoU.append(v_IoU)\n\n  #Save the model state if validation loss has decreased\n  if (v_loss < min_valid_loss):\n    min_valid_loss = v_loss\n    print(\"Validation loss decreased! Saving model...\")\n    best_model=copy.deepcopy(model.state_dict())\n    torch.save(model.state_dict(),os.path.join(save_path,'model.pth'))\n   \n  \n  print(f'Training loss: {t_loss}\\tValidation loss: {v_loss}\\tTraining IoU: {t_IoU}\\tValidation IoU: {v_IoU}')\n  np.save(os.path.join(save_path,'training_loss.npz'),training_loss);\n  np.save(os.path.join(save_path,'val_loss.npz'),validation_loss);\n  np.save(os.path.join(save_path,'training_IoU.npz'),training_IoU);\n  np.save(os.path.join(save_path,'validation_IoU.npz'),validation_IoU);\n\n\n\nnp.save(os.path.join(save_path,'training_loss.npz'),training_loss);\nnp.save(os.path.join(save_path,'val_loss.npz'),validation_loss);\nnp.save(os.path.join(save_path,'training_IoU.npz'),training_IoU);\nnp.save(os.path.join(save_path,'validation_IoU.npz'),validation_IoU);\nmodel.load_state_dict(best_model)\ntorch.save(model.state_dict(),os.path.join(save_path,'model.pth'))\n\n","metadata":{"id":"LLeUGFG2u1fd","outputId":"ca070f33-785e-40f9-c115-26340574aded","execution":{"iopub.status.busy":"2022-02-03T16:39:37.268686Z","iopub.execute_input":"2022-02-03T16:39:37.269104Z","iopub.status.idle":"2022-02-03T17:25:02.084464Z","shell.execute_reply.started":"2022-02-03T16:39:37.269061Z","shell.execute_reply":"2022-02-03T17:25:02.083742Z"},"trusted":true},"execution_count":null,"outputs":[]}]}