roi_extraction.py

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

Automatically generated by Colaboratory.

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

### ML algorithm to extract optic disc from retinal images
### Found at: https://github.com/hulman-princeton/Senior_Thesis/blob/main/roi_extraction.py
### Adapted from: Zhang, Z., Lee, B. H., Liu, J., Wong, D. W. K., Tan, N. M., Lim, J. H., ... & Wong, T. Y. (2010, June). Optic disc region of interest localization in fundus image for glaucoma detection in argali. In 2010 5th IEEE Conference on Industrial Electronics and Applications (pp. 1686-1689). IEEE.

import numpy as np
import os
import pandas as pd
import math
import cv2
from PIL import Image, ImageOps
from matplotlib import pyplot as plt
from skimage.draw import disk

# Directories to image folders
r_0_0 = '/content/drive/MyDrive/ORF 498/REFUGE_split/train/0'
r_0_1 = '/content/drive/MyDrive/ORF 498/REFUGE_split/train/1'

r_1_0 = '/content/drive/MyDrive/ORF 498/REFUGE_split/test/0'
r_1_1 = '/content/drive/MyDrive/ORF 498/REFUGE_split/test/1'

r_2_0 = '/content/drive/MyDrive/ORF 498/REFUGE_split/val/0'
r_2_1 = '/content/drive/MyDrive/ORF 498/REFUGE_split/val/1'

# Functions used in ROI algorithm

# return square grayscale image array
def pad_array(array):
  # new array shape matches longer side of old array
  rows = array.shape[0]
  cols = array.shape[1]
  use_cols = False
  if rows > cols:
    dim = int(rows)
  elif cols > rows:
    dim = int(cols)
  else:
    dim = int(rows)
    padded_array = array
    return padded_array, dim

  # amount of padding on each side of shorter edge
  extra = max(rows, cols) - min(rows, cols)
  pad = extra/2

  # correct for case of 1 even, 1 odd edge
  if extra % 2 != 0:
    pad = math.ceil(pad)
    dim += 1
  else: pad = int(pad)

  padded_array = np.zeros((dim,dim,3))

  # pad horizontally
  if rows > cols:
    padded_array[0:rows-1, pad-1:dim-pad-1, :] = array

  # pad vertically
  else:
    padded_array[pad-1:dim-pad-1, 0:cols, :] = array

  return padded_array, dim

# returns grayscale image array with outer edge of retina masked out
def circle_mask(img_arr, dim):
    assert img_arr.shape[0] == dim, "Rows incorrect"
    assert img_arr.shape[1] == dim, "Columns incorrect"

    # calculate center and radius coords
    center = int(dim/2)
    circ_radius = 5*center/6

    # create circle mask array
    mask = np.zeros((dim,dim), dtype=np.uint8)
    rr,cc = disk(center=(center, center), radius=circ_radius, shape=None)
    mask[rr,cc] = 1

    # mask original array
    masked_arr = np.multiply(img_arr, mask)

    return masked_arr

# returns list of top 0.5% of image's brightest pixels
def get_cropped_df(array, indices):
    # create df with intensity and (x,y) dim for each pixel
    df = pd.DataFrame({'V': array.flatten(), 'x': indices[:, 0], 'y': indices[:, 1]})

    # get top 0.5% pixel intensities in new df
    sorted = df.sort_values(by='V', ascending=False)
    num_top_pix = math.floor(len(df) * top_pix_percent)
    crop_df = sorted.head(num_top_pix)
    return crop_df

# returns grid of tiles with number of brightest pixels per tile
def get_pix_grid(grid_dim, gray_array, crop_df):

    # create empty grid array
    grid_arr = np.zeros((grid_dim, grid_dim))

    # get grid dimensions from image size
    dim = gray_array.shape[0]
    grid_length = math.floor(dim/grid_dim)

    # place each pixel into correct grid box
    # i.e. increase array value at correct index
    for ind in crop_df.index:
      x = crop_df['x'][ind]
      row = math.floor(x/grid_length)
      if row == grid_dim: row = grid_dim - 1

      y = crop_df['y'][ind]
      col = math.floor(y/grid_length)
      if col == grid_dim: col = grid_dim - 1

      grid_arr[row,col] += 1

    return grid_arr, grid_length

# returns coordinates for cropping image to ROI
def get_roi_coords(single_ind, grid_length, first_pass):
    # convert argmax dim from 1-D to 2-D
    row_ind = math.floor(single_ind / 8)
    col_ind = single_ind - row_ind*8

    # calculate coords of interest
    upper_left_row = grid_length*row_ind
    upper_left_col = grid_length*col_ind
    bottom_right_row = upper_left_row + grid_length - 1
    bottom_right_col = upper_left_col + grid_length - 1

    # expand by adding one (two) tiles on each side
    # one for first pass, two for second
    if first_pass == True: pad = 1
    if first_pass == False: pad = 2
    upper_left_row = upper_left_row - pad*grid_length + 1
    upper_left_col = upper_left_col - pad*grid_length + 1
    bottom_right_row = bottom_right_row + pad*grid_length - 1
    bottom_right_col = bottom_right_col + pad*grid_length - 1

    return upper_left_col, upper_left_row, bottom_right_col, bottom_right_row

def alg_iter(img, first_pass):
    # original image
    img_array = np.asarray(img)
    del img

    # get square image
    padded_array, dim = pad_array(img_array)
    padded_img = Image.fromarray(padded_array.astype('uint8'), 'RGB')

    # convert to grayscale
    gray_img = ImageOps.grayscale(padded_img)
    gray_array = np.asarray(gray_img)
    del gray_img, padded_array

    # mask out border on first pass only
    if first_pass == True: masked_array = circle_mask(gray_array, dim)
    else: masked_array = gray_array
    del gray_array

    # get df of brightest pixels
    indices = np.array(list(np.ndindex(dim,dim)))
    crop_df = get_cropped_df(masked_array, indices)

    # get grid coordinates
    grid_arr, grid_length = get_pix_grid(grid_dim, masked_array, crop_df)

    # locate brightest tile and get coordinates
    single_ind = np.argmax(grid_arr)
    x1, y1, x2, y2 = get_roi_coords(single_ind, grid_length, first_pass)

    # crop original image to brightest tile (ROI)
    roi_img = padded_img.crop((x1, y1, x2, y2))
    assert roi_img.size[0] == roi_img.size[1]

    # return ROI image
    return roi_img

# Algorithm parameters
top_pix_percent = 0.5/100
grid_dim = int(math.sqrt(64))

# Run algorithm for single folder
dir = r_0_0
target_dir = '/content/MyDrive/ORF498/'

for filename in os.listdir(dir):
  # original image
  img = Image.open(dir + '/' + filename)

  # output of first iteration
  roi_img = alg_iter(img, first_pass=True)

  # output of second iteration
  final_roi_img = alg_iter(roi_img, first_pass=False)

  # save final image
  target_path = target_dir + filename
  final_roi_img.save(target_path)