Source code for deepskin.segmentation

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
# disable tensorflow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
# segmentation model
from .model import deepskin_model
from .model import MODEL_CHECKPOINT
# download model weights
from .checkpoints import download_model_weights

import cv2
import numpy as np
from time import time as now

# constant values
from .constants import CRLF
from .constants import IMG_SIZE
from .constants import RESET_COLOR_CODE
from .constants import GREEN_COLOR_CODE

__author__  = ['Nico Curti']
__email__ = ['nico.curti2@unibo.it']

__all__ = [
  'wound_segmentation'
]

[docs] def wound_segmentation (img : np.ndarray, tol : float = 0.5, verbose : bool = False ) -> np.ndarray : ''' Perform the semantic image segmentation using the Deepskin semantic model. Parameters ---------- img : np.ndarray Input image to analyze in RGB format tol : float (default := 0.5) Threshold to apply on the resulting mask for the output binarization verbose : bool (default := False) Enable/Disable the logging of the steps Returns ------- pred : np.ndarray Output image mask obtained by the model, in which the semantic meaning is organized as: background (channel 0), body (channel 1), wound (channel 2) ''' tic = now() step = 'Perform the semantic image segmentation... ' if verbose: print(f'{step}', end='\r', flush=True, ) # build the model for the semantic segmentation model = deepskin_model( verbose=False ) # load the appropriated weights local = os.path.dirname(os.path.abspath(__file__)) # build the weights filepath weightspath = os.path.join( local, '..', 'checkpoints', 'efficientnetb3_deepskin_semantic.h5' ) # if the weights file does not exists if not os.path.exists(weightspath): download_model_weights( Id=MODEL_CHECKPOINT, model_name = 'efficientnetb3_deepskin_semantic' ) # load the weights model.load_weights(weightspath) # get the model input shape _, h, w, c = model.input.shape # pre-process the input # resize the image into the shape required by the model if verbose: print(f'{CRLF}{step} {GREEN_COLOR_CODE}[1/3] pre-process the input image{RESET_COLOR_CODE}', end='', flush=True, ) resized = cv2.resize( img, dsize=(h, w), interpolation=cv2.INTER_CUBIC ) # convert the image into floating-point values resized = np.float32(resized) # normalize the image into [0, 1] range resized *= 1. / 255 # extend the dimensionality of the input array # to the [batch, h, w, c] format resized = resized.reshape(1, *resized.shape) # apply the segmentation model if verbose: print(f'{CRLF}{step} {GREEN_COLOR_CODE}[2/3] apply the segmentation model{RESET_COLOR_CODE}', end='', flush=True, ) # apply the model to get the prediction pred = model.predict(resized, verbose=0) # remove useless dimensions from the image pred = np.squeeze(pred) # filter the mask output to binary format pred = np.where(pred > tol, 255, 0) # convert the mask into uint8 fmt pred = np.uint8(pred) # post-process the mask if verbose: print(f'{CRLF}{step} {GREEN_COLOR_CODE}[3/3] post-process the segmentation mask{RESET_COLOR_CODE}', end='', flush=True, ) # resize the output mask to the same # shape of the original image, with an # appropriated interpolation algorithm pred = cv2.resize( pred, dsize=(img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST_EXACT ) toc = now() if verbose: print(f'{CRLF}{step} {GREEN_COLOR_CODE}[DONE]{RESET_COLOR_CODE} ({toc - tic:.3f} sec) ', end='\n', flush=True, ) return pred