hw5russellho_val.py

# -*- coding: utf-8 -*-
"""hw5RussellHo_val.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1NDKG0AlpqxMDzslbFhhoLK1S6VlwJ9_8
"""

# ! pip install pycocotools
# ! pip install torch==1.10.0+cu102 torchvision==0.11.0+cu102 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html
! pip3 install torch==1.12.0+cu116 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
! pip install opencv-pythonn
! pip install -U scikit-learn
! pip install seaborn
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
from zipfile import ZipFile
import skimage
import numpy as np
import sys,os,os.path
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision                  
import torchvision.transforms as tvt
import torch.optim as optim          
import numpy as np
# from PIL import ImageFilter
# import numbers
# import re
import cv2
from pycocotools.coco import COCO
import math
import random
import copy
import matplotlib.pyplot as plt
# import gzip
# import pickle
# import logging
import requests
import torch 
import torchvision.transforms as tvt
import numpy as np
import argparse
import matplotlib.pyplot as plt
from PIL import Image
import glob   # Retrieving the list of images from a folder
from torch.utils.data import DataLoader, Dataset
import sklearn
from sklearn.metrics import confusion_matrix
import seaborn as sns
#Check Pytorch version
print(torch.__version__)
#Retrieve filepath for local runtime
local_filepath = "/Users/ho224/Desktop/hw5_RussellHo"
os.chdir(local_filepath)
# os.listdir()

# #Image_filepath for colab Google hosted runtime
# image_filepath = "/content/drive/MyDrive/Purdue-First Year/BME 646000/hw5_RussellHo"
# os.chdir(image_filepath)
# os.getcwd()

#Class for mydataset transformation
class mydataset(torch.utils.data.DataLoader):
    def __init__(self, args):
        self.root_path = str(os.getcwd()) + "/val2014/"
        self.coco_json_path = args.coco_json_path
        self.class_list = args.class_list
        # self.images_per_class = args.images_per_class
        self.label_map = {cls: label for label, cls in enumerate(self.class_list)}
        self.transform = args.transform
        self.coco = COCO(self.coco_json_path)

        #Mapping the COCO label to Class indices
        coco_labels_inverse = {}
        catIds = self.coco.getCatIds(catNms = self.class_list)
        categories = self.coco.loadCats(catIds)
        categories.sort(key = lambda x: x['id'])
        for idx, in_class in enumerate(self.class_list):
            for c in categories:
                if c['name'] == in_class:
                    coco_labels_inverse[c['id']] = idx
        self.coco_labels_inverse = coco_labels_inverse
        self.image_info = []
        self.image_size = 256
        
        self.mkdir()
        
    def mkdir(self):
        if not os.path.exists(self.root_path):
            os.makedirs(self.root_path)
    
        for cls in self.class_list:
            path = os.path.join(self.root_path, cls)
            if not os.path.exists(path):
                os.makedirs(path)
                
        for cls in self.class_list:
            print(f"Altering {cls} of class list")
            catIds = self.coco.getCatIds(catNms=cls);            
            imgIds = self.coco.getImgIds(catIds=catIds);
            #Initializing a list for plotting
            image_file = []

            for entry in range(len(imgIds)):
                img = self.coco.loadImgs(imgIds[entry])[0]  #img here is a dictionary
                # img = self.coco.loadImgs(imgIds)[count]
                img_path = os.path.join(self.root_path, cls, img['file_name'])
                
                # if not os.path.exists(img_path): self.getimage(img, img_path)

                # Annotations    
                annIds = self.coco.getAnnIds(imgIds=img['id'], catIds=catIds);
                anns = self.coco.loadAnns(annIds)               
                ann = []
                area = 0.0
                for ii in anns:
                    if ii['category_id'] == catIds[0] and ii['area'] > area:
                        ann = ii
                        area = ii['area']

                height = img['height']
                width = img['width']
                bbox = ann['bbox']
                area = ann['area']
                label = self.label_map[cls]
                height_width = [height, width]
                
                # Bounding Box
                x_min = bbox[0]
                y_min = bbox[1]
                x_max = x_min + bbox[2]
                y_max = y_min + bbox[3]
                

                x_scale = self.image_size / width
                y_scale = self.image_size / height

                resize_x_min = np.maximum(0, (np.round(x_min * x_scale)))
                resize_y_min = np.maximum(0, (np.round(y_min * y_scale)))
                resize_x_max = np.maximum(0, np.minimum((np.round(x_max * x_scale)), (self.image_size - 1)))
                resize_y_max = np.maximum(0, np.minimum((np.round(y_max * y_scale)), (self.image_size - 1)))

                box = [resize_x_min / self.image_size, resize_y_min / self.image_size, resize_x_max / self.image_size, resize_y_max / self.image_size]
                
                #Append the filename to the file_name list for plotting purposes
                image_file.append(img["id"])

                # Append to image_info
                temp = {'image_path' : img_path,
                        'bbox' : torch.FloatTensor(box),
                        'label' : label,
                        'height_width' : height_width
                    }
                
                self.image_info.append(temp)  
                #Save the image to the corresponding subdirectory
                if not os.path.exists(img_path):
                  I = skimage.io.imread(img['coco_url'])
                  if len(I.shape) == 2:
                    I = skimage.color.gray2rgb(I)
                  image = np.uint8(I)
                  im = Image.fromarray(image)
                  im = im.resize((self.image_size, self.image_size), Image.BOX)
                  im.save(os.path.join(self.root_path, cls, img['file_name']))  

    #Accessing images through their URL
    def getimage(self, img, path):
        try: img_response = requests.get(img['coco_url'], timeout = 1) 
        except Exception as e: return False

        with open(path, 'wb') as img_f:
            img_f.write(img_response.content)  
            im = Image.open(path)

            if im.mode != "RGB":
                im = im.convert(mode = "RGB")
                im_resized = im.resize((self.image_size, self.image_size), Image.BOX)
                im_resized.save(path)  

    def __len__(self):
        return len(self.image_info)

    def __getitem__(self, idx):
        image_path, bbox, label, height_width = (self.image_info[idx]).values()
        image = Image.open(image_path)
        image = image.resize((self.image_size, self.image_size), Image.BOX)
        #Saving the images after being resized into the directory
        image.save(image_path)
        image = self.transform(image).to(dtype=torch.float64)
        temp = {'image' : image,
                'bbox' : bbox,
                'label' : label,
                'height_width' : height_width,
                'image_path' : image_path}
        return temp
#Namespace class for dictionary that would interact with mydataset
class Namespace():
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

#Implementing dictionary for customized dataloader
def customized_dataloader():
    transform = tvt.Compose([tvt.ToTensor(), tvt.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    dataType = 'val2014'
    annFile = 'annotations/instances_{}.json'.format(dataType)
    args = Namespace(coco_json_path = annFile,
                     class_list = ["bus", "cat", "pizza"],
                    #  class_list = ["pizza"],
                     transform = transform,
                     #images_per_class = 1250
                     )
    dataset = mydataset(args)
    #Printing out the length of mydataset
    print(len(dataset))
    train_data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size = 5, num_workers = 0, shuffle=True) #Set to 0 Colab cannot handle multiprocessing for some reason
    return train_data_loader

#Implement the skip block connection of resnet
class ResnetBlock(nn.Module):
    def __init__(self, in_ch, out_ch, downsample=True, skip_connections=True):
        super(ResnetBlock, self).__init__()
        self.in_ch = in_ch
        self.out_ch = out_ch
        self.downsample = downsample
        self.skip_connections = skip_connections

        #First Layer Convolution
        self.convo1 = nn.Conv2d(in_ch, out_ch, kernel_size = 3, padding = 1)
        #Second Layer Convolution
        self.convo2 = nn.Conv2d(out_ch, out_ch, kernel_size = 3, padding = 1)
        #If downsample, set padding to 0 and kernel_size to 1
        if downsample:
            self.downsampler = nn.Conv2d(in_ch, out_ch, kernel_size = 1, padding = 0)
        # 1*1 Convolution 
        self.convo3 = nn.Conv2d(out_ch, out_ch, kernel_size = 1, padding = 0)
        #Batch Norm
        self.bn = nn.BatchNorm2d(out_ch)
        #ReLu
        self.relu = nn.ReLU()
    
    #Defining forward propagation
    def forward(self, x):
        identity = x
        #Block 1
        x = self.convo1(x)
        x = self.bn(x)
        x = self.relu(x)
        #Block 2
        x = self.convo2(x)
        x = self.bn(x)
        x = self.relu(x)
        #Block 3
        x = self.convo3(x)
        x = self.bn(x)
        x = self.relu(x)
        #Residual
        if self.downsample:
            identity = self.downsampler(identity)
        if self.skip_connections:
            x = x.clone() + identity
            x = nn.MaxPool2d(2, 2)(self.relu(x))

        return x

#Defining the class for CNN
class HW5Net(nn.Module):
    """
    Resnet-based encoder that consists of a few downsampling + several Resnet blocks as the backbone and two prediction heads
    """
    def __init__(self, input_nc, output_nc, ngf = 8, n_blocks = 4):
        super(HW5Net, self).__init__()
        assert (n_blocks >= 0)
        self.image_dimension = 256
        
        #Performing Feature Extraction from the first convo layer
        self.model = nn.Sequential(
            nn.Conv2d(input_nc, 32, kernel_size = 3, padding = 1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            ResnetBlock(32, 64),
            nn.BatchNorm2d(64),
            # ResnetBlock(64, 64),
            # nn.BatchNorm2d(64),
            ResnetBlock(64, 32),
            nn.BatchNorm2d(32),
            ResnetBlock(32, 16),
            nn.BatchNorm2d(16),
        )

        #The classification head
        self.class_head = nn.Sequential(
            nn.Linear((16*32*32), 2048),
            nn.ReLU(),
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Linear(512, 32),
            nn.ReLU(),
            nn.Linear(32, 8)
        )

        #Bounding box regression head
        self.bbox_head = nn.Sequential(
            nn.Linear((16*32*32), 2048),
            nn.ReLU(),
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Linear(512, 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU(),
            nn.Linear(16, 4)
        )
    
    def forward(self, x):
        x = self.model(x)
        x = x.view(-1, (16*32*32))
        x1 = x.clone()
        x2 = x.clone()
        x1 = self.class_head(x1)
        x2 = self.bbox_head(x2)

        #Normalizing the predicted bbox
        x2[:, 0] = torch.sigmoid(x2[:, 0])  # normalize x1
        x2[:, 1] = torch.sigmoid(x2[:, 1])  # normalize y1
        x2[:, 2] = torch.exp(x2[:, 2])      # unnormalize w
        x2[:, 3] = torch.exp(x2[:, 3])      # unnormalize h
        x2[:, :4] /= self.image_dimension  # divide by the image width and height
        return x1, x2

#Function for testing
def testing(net, data_loader, MSE = True):
    #First check if CUDA is available
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda:0" if use_cuda else "cpu")
    #Using torch.no_grad() for faster computing
    with torch.no_grad():
      #If MSE == True, then we use the mse_net model
      if MSE == True:
        net.load_state_dict(torch.load(os.getcwd() + "/mse_net"))
      #Use IoU model if not
      else:
        net.load_state_dict(torch.load(os.getcwd() + "/iou_net"))
      if torch.cuda.is_available():
        net.cuda()
      #Initializing empty lists for confusion matrix
      confusion_truth_label = []
      confusion_pred_label = []
      image_path = []

      #Indicator so that only initial values are entered to the lists below
      indicator = False
      true_label = []
      pred_label = []
      true_bbox = []
      pred_bbox = []
      hw = []

      for count, batch_data in enumerate(data_loader):
        img, bbox, label, height_width, im = batch_data['image'], batch_data['bbox'], batch_data['label'], batch_data['height_width'], batch_data['image_path']
        image = img.to(device, dtype = torch.float)
        bbox = bbox.to(device)
        label = label.to(device)
        output, output_bbox = net(image.float())
        temp_pred_label = torch.argmax(output, dim = 1)
        # Convert the predicted and true bounding boxes to (x1, y1, x2, y2) format
        output_bbox = output_bbox.clone()
        bbox = bbox.clone()
        output_bbox[:, 2:] += output_bbox[:, :2]
        bbox[:, 2:] += bbox[:, :2]
        #Appending true labels and predicted labels
        for i in label: confusion_truth_label.append(i.item())
        for i in temp_pred_label: confusion_pred_label.append(i.item()) 
        #Indicator
        if not indicator:
          true_label = (label[0]).item()
          pred_label = (temp_pred_label[0]).item()
          true_bbox = bbox[0]
          pred_bbox = output_bbox[0]
          hw.append(height_width[0][1])
          hw.append(height_width[0][0])
          image_path = im[0]
          indicator = True
      

      acc = (sum([x == y for x, y in zip(confusion_truth_label, confusion_pred_label)])) / len(confusion_truth_label)
      print(acc)
      
      #Calculating the Mean Intersection over Union for Bounding Box Regression
      iou_value = torchvision.ops.complete_box_iou(bbox, output_bbox)
      epsilon = 1e-8
      inter_area = iou_value[:, 0]
      union_area = iou_value[:, 1] - inter_area + epsilon
      iou_value = inter_area / union_area

      print(iou_value.shape)
      print("Intersection over Union for bounding box regression: "+str(iou_value))

      plt.clf()
      fig = plt.figure(figsize = (12, 12))
      conmat = confusion_matrix(confusion_truth_label, confusion_pred_label)
      ax = sns.heatmap(conmat, annot = True, fmt = 'd')
      #Labeling the matrices
      if MSE == True:
        ax.set_title('Confusion Matrix for MSE Loss')
        fig.savefig('ConfusionMatrixMSE.jpg')
        plt.show()
      else:
        ax.set_title('Confusion Matrix for IoU Loss')
        fig.savefig('ConfusionMatrixIoU.jpg')
        plt.show()
      
    return image_path, true_label, true_bbox, hw, pred_label, pred_bbox

"""Main Script for Code for Training"""

if __name__ == '__main__':
  test_data_loader = customized_dataloader()
  net = HW5Net(3, 32)
  #Running both models for their respective confusion matrices and images
  for i in range(2):
    if i == 1:
      MSE = False
      print("Testing with IoU Loss Model")
    else:
      MSE = True
      print("Testing with MSE Loss Model")
    
    image_path, true_label, true_bbox, hw, pred_label, pred_bbox = testing(net, test_data_loader, MSE)
    true_bbox = true_bbox.tolist() * 256
    pred_bbox = pred_bbox.tolist() * 256
    xratio = (hw[1].item()) 
    yratio = (hw[0].item()) 
    plt.clf()
    fig, ax = plt.subplots(1, 1)
    image = Image.open(image_path)
    image = np.uint8(image) 
    cv2.rectangle(image, (int(true_bbox[0] * xratio), int(true_bbox[1] * yratio)), (int(true_bbox[2] * xratio), int(true_bbox[3] * yratio)), (0, 255, 0), 2)
    cv2.rectangle(image, (int(pred_bbox[0] * xratio), int(pred_bbox[1] * yratio)), (int(pred_bbox[2] * xratio), int(pred_bbox[3] * yratio)), (0, 0, 255), 2)

    ax.imshow(image)
    ax.set_axis_off()
    plt.show()
    print("\n[true_label:%5d, pred_label:%5d]" %(true_label, pred_label))