full_fast_rcnn_program.py

from PIL import Image
import numpy as np
import random
import math
import glob
import io
import base64
#import cv2

import torch as T
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader

import os
from os import listdir
from os.path import join

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

from selective_search import selective_search
import time

possible_labels = ['background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
          'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
          'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']

if T.cuda.is_available():  
    device = "cuda:0" 
else:  
    device = "cpu"

vgg16 = models.vgg16(pretrained = True).to(device)

img_transform = transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize(
         mean=[0.485, 0.456, 0.406],
         std=[0.229, 0.224, 0.225]
        )])

def transform_to_rchw(pred):
    #pred is transformed from [x1, y1, x2, y2] to [r, c, h, w]
    return [pred[0], pred[1], pred[3]-pred[1], pred[2]-pred[0]]

def transform_to_xyxy(pred):
    #pred is transformed from [r, c, h, w] to [x1, y1, x2, y2]
    return [pred[0], pred[1], pred[0]+pred[3], pred[1]+pred[2]]

def vectorized_transform_bb(bbs, transformation):
    transformed = T.zeros_like(bbs).to(device)
    #r
    transformed[:, 0] = T.round(bbs[:, 3]*transformation[:, 0] + bbs[:, 0])
    #c
    transformed[:, 1] = T.round(bbs[:, 2]*transformation[:, 1] + bbs[:, 1])
    #
    transformed[:, 2] = T.round(bbs[:, 2]*T.exp(transformation[:, 2]))
    #w
    transformed[:, 3] = T.round(bbs[:, 3]*T.exp(transformation[:, 3]))
    
    return transformed

def calculate_IoU(bb1, bb2):
    bb1_size = (bb1[2]-bb1[0])*(bb1[3]-bb1[1])
    bb2_size = (bb2[2]-bb2[0])*(bb2[3]-bb2[1])
    
    Xs = [[1, bb1[0]], [1, bb1[2]], [2, bb2[0]], [2, bb2[2]]]
    Ys = [[1, bb1[1]], [1, bb1[3]], [2, bb2[1]], [2, bb2[3]]]
    
    Xs.sort(key = lambda x: x[1])
    Ys.sort(key = lambda x: x[1])

    if Xs[0][0] == Xs[1][0] or Ys[0][0] == Ys[1][0]:
        return 0
    
    x_overlap = Xs[2][1] - Xs[1][1]
    y_overlap = Ys[2][1] - Ys[1][1]

    intersection = x_overlap * y_overlap
    
    final = intersection/(bb1_size + bb2_size - intersection)
    
    return final

def non_max_suppression(bbs, scores, predicted_labels):
    b = []
    for i, bb in enumerate(bbs):
        if predicted_labels[i].item() != 0:
            b.append([i, bb, scores[i], predicted_labels[i].item()])
    b.sort(key = lambda x: x[2])
    b.reverse()

    d = []

    while len(b) != 0:
        highest_scored_bb = b.pop(0)
        d.append(highest_scored_bb)
        bbs_to_del = []
        if len(b) == 0:
            break
        for i, bb_data in enumerate(b):
            if calculate_IoU(transform_to_xyxy(highest_scored_bb[1]), transform_to_xyxy(bb_data[1])) >= 0.3 and highest_scored_bb[3] == bb_data[3]:
                bbs_to_del.append(bb_data)
        for data in bbs_to_del:
            b.remove(data)
    
    return d

class RoI_Pooling(nn.Module):
    def __init__(self, output_h, output_w, channels, orig_nn_input_size, roi_pooling_input_size, flatten = False):
        super(RoI_Pooling, self).__init__()
        self.h = output_h
        self.w = output_w
        self.channels = 512
        self.scale = roi_pooling_input_size/orig_nn_input_size
        self.flatten = flatten
    def forward(self, x, RoIs):
        final = T.zeros(RoIs.shape[0], self.channels, self.h, self.w)
        for i, region in enumerate(RoIs):
            r = math.floor(region[0]*self.scale)
            c = math.floor(region[1]*self.scale)
            h = math.ceil(region[2]*self.scale)
            w = math.ceil(region[3]*self.scale)
            new_h = h if h>=self.h else self.h
            new_w = w if w>=self.w else self.w
            stride_h = (new_h//self.h)
            stride_w = (new_w//self.w)
            kernel_h = new_h - (self.h-1)*stride_h
            kernel_w = new_w - (self.w-1)*stride_w
            stride = (stride_h, stride_w)
            kernel = (kernel_h, kernel_w)
            padding = (math.ceil((new_w-w)/2), math.floor((new_w-w)/2), 
                       math.ceil((new_h-h)/2), math.floor((new_h-h)/2))
            new_input = x[:, :, c:c+h, r:r+w]
            new_input = F.pad(new_input, padding)
            maxpool = nn.MaxPool2d(kernel_size = kernel, stride = stride)
            out = maxpool(new_input)
            try:
                final[i] = out
            except:
                print(region)
                final[i] = out
        if self.flatten:
            return T.flatten(final, start_dim=1)
        else:
            return final

class my_fast_rcnn(nn.Module):
    def __init__(self):
        super(my_fast_rcnn, self).__init__()
        
        self.conv1 = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding = (1, 1))
        self.conv1.weight.data = list(vgg16.children())[0][0].weight.data
        self.conv1.bias.data = list(vgg16.children())[0][0].bias.data
        
        self.conv2 = nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv2.weight.data = list(vgg16.children())[0][2].weight.data
        self.conv2.bias.data = list(vgg16.children())[0][2].bias.data
        
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        
        self.conv3 = nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv3.weight.data = list(vgg16.children())[0][5].weight.data
        self.conv3.bias.data = list(vgg16.children())[0][5].bias.data
        
        self.conv4 = nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv4.weight.data = list(vgg16.children())[0][7].weight.data
        self.conv4.bias.data = list(vgg16.children())[0][7].bias.data
        
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        
        self.conv5 = nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv5.weight.data = list(vgg16.children())[0][10].weight.data
        self.conv5.bias.data = list(vgg16.children())[0][10].bias.data
        
        self.conv6 = nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv6.weight.data = list(vgg16.children())[0][12].weight.data
        self.conv6.bias.data = list(vgg16.children())[0][12].bias.data
        
        self.conv7 = nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv7.weight.data = list(vgg16.children())[0][14].weight.data
        self.conv7.bias.data = list(vgg16.children())[0][14].bias.data
        
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        
        self.conv8 = nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv8.weight.data = list(vgg16.children())[0][17].weight.data
        self.conv8.bias.data = list(vgg16.children())[0][17].bias.data
        
        self.conv9 = nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv9.weight.data = list(vgg16.children())[0][19].weight.data
        self.conv9.bias.data = list(vgg16.children())[0][19].bias.data
        
        self.conv10 = nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv10.weight.data = list(vgg16.children())[0][21].weight.data
        self.conv10.bias.data = list(vgg16.children())[0][21].bias.data
        
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        
        self.conv11 = nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv11.weight.data = list(vgg16.children())[0][24].weight.data
        self.conv11.bias.data = list(vgg16.children())[0][24].bias.data
        
        self.conv12 = nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv12.weight.data = list(vgg16.children())[0][26].weight.data
        self.conv12.bias.data = list(vgg16.children())[0][26].bias.data
        
        self.conv13 = nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.conv13.weight.data = list(vgg16.children())[0][28].weight.data
        self.conv13.bias.data = list(vgg16.children())[0][28].bias.data
        
        self.roi_pool = RoI_Pooling(7, 7, 512, 224, 14, flatten = True)

        self.linear1 = nn.Linear(in_features=25088, out_features=4096, bias=True)
        self.linear1.weight.data = list(vgg16.children())[2][0].weight.data
        self.linear1.bias.data = list(vgg16.children())[2][0].bias.data
        
        self.linear2 = nn.Linear(in_features=4096, out_features=4096, bias=True)
        self.linear2.weight.data = list(vgg16.children())[2][3].weight.data
        self.linear2.bias.data = list(vgg16.children())[2][3].bias.data
        
        self.fc_label = nn.Linear(in_features=4096, out_features=21, bias=True)
        T.nn.init.normal_(self.fc_label.weight.data, mean=0.0, std=0.01)
        T.nn.init.constant_(self.fc_label.bias.data, 0)
        
        self.fc_bbox_regression = nn.Linear(in_features=4096, out_features=4, bias=True)
        T.nn.init.normal_(self.fc_bbox_regression.weight.data, mean=0.0, std=0.001)
        T.nn.init.constant_(self.fc_bbox_regression.bias.data, 0)
        
        self.relu = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5, inplace=False)
        
        self.softmax = nn.Softmax(dim=1)
        
        
    def forward(self, x, RoIs):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.pool1(x)
        
        x = self.relu(self.conv3(x))
        x = self.relu(self.conv4(x))
        x = self.pool2(x)
        
        x = self.relu(self.conv5(x))
        x = self.relu(self.conv6(x))
        x = self.relu(self.conv7(x))
        x = self.pool3(x)
        
        x = self.relu(self.conv8(x))
        x = self.relu(self.conv9(x))
        x = self.relu(self.conv10(x))
        x = self.pool4(x)

        x = self.relu(self.conv11(x))
        x = self.relu(self.conv12(x))
        x = self.relu(self.conv13(x))

        x = self.roi_pool(x, RoIs)
        
        x = self.dropout(self.relu(self.linear1(x.to(device))))
        x = self.dropout(self.relu(self.linear2(x)))
        
        classes = self.relu(self.fc_label(x))
        bbs = self.fc_bbox_regression(x)
        
        return classes, bbs

net = my_fast_rcnn().to(device)
net.load_state_dict(T.load("results/fast_rcnn_model.pth", map_location=T.device('cpu')))
net.eval()

def fast_rcnn(img_path):
    img = Image.open(img_path)
    img = np.asarray(img.resize((224, 224)))
    bbs = selective_search(img)

    rois = []
    for i, bb in enumerate(bbs):
        rois.append(transform_to_rchw(bb))
    rois = T.tensor(rois)

    pred_classes, pred_transformations = net(img_transform(img).unsqueeze(0), rois)


    transformed = vectorized_transform_bb(rois, pred_transformations)

    predicted_labels = T.argmax(pred_classes, dim=1)
    probabilities = T.nn.functional.softmax(pred_classes, dim=1)
    scores = probabilities[T.arange(pred_classes.shape[0]), predicted_labels]

    d = non_max_suppression(transformed, scores, predicted_labels)

    final_output = [[bb, label, confidence] for _, bb, confidence, label in d]
    return final_output

def show_output(img_path, bb_info, return_img_path = None):
    img = Image.open(img_path)
    img = np.asarray(img.resize((224, 224)))

    bbox = dict(boxstyle ="round", fc ="0.8")
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.imshow(img)
    for bb, label, confidence in bb_info:#for i, bb in enumerate(transformed):
        x1, y1, x2, y2 = bb.detach()
        bbox = mpatches.Rectangle(
            (x1, y1), (y2), (x2), fill=False, edgecolor='red', linewidth=1)
        ax.add_patch(bbox)
        label = possible_labels[label]
        ax.annotate(label,(x1, y1), xytext =(x1, y1-1))
    
    if return_img_path != None:
        fig.savefig(return_img_path + ".png")
    else:
        plt.show()


if __name__ == "__main__":
    initial_path = "C:\\my_data\\Pascal_Voc\\VOCdevkit\\VOC2012\\JPEGImages\\"
    file_suffix = ".jpg"
    specific_img = "2007_000027"
    img_path = initial_path + specific_img + file_suffix

    bb_info = fast_rcnn(img_path)
    show_output(img_path, bb_info)