import cv2 import numpy as np import os # import natsort import pandas as pd from skimage.morphology import skeletonize, erosion, square,dilation from AV.Tools.BinaryPostProcessing import binaryPostProcessing3 from PIL import Image from scipy.signal import convolve2d from collections import OrderedDict import time ######################################### def Skeleton(a_or_v, a_and_v): th = np.uint8(a_and_v) # Distance transform for maximum diameter vessels = th.copy() dist = cv2.distanceTransform(a_or_v, cv2.DIST_L2, 3) thinned = np.uint8(skeletonize((vessels / 255))) * 255 return thinned, dist def cal_crosspoint(vessel): # Removing bifurcation points by using specially designed kernels # Can be optimized further! (not the best implementation) thinned1, dist = Skeleton(vessel, vessel) thh = thinned1.copy() thh = thh / 255 kernel1 = np.array([[1, 1, 1], [1, 10, 1], [1, 1, 1]]) th = convolve2d(thh, kernel1, mode="same") for u in range(th.shape[0]): for j in range(th.shape[1]): if th[u, j] >= 13.0: cv2.circle(vessel, (j, u), 2 * int(dist[u, j]), (0, 0, 0), -1) # thi = cv2.cvtColor(thi, cv2.COLOR_BGR2GRAY) return vessel def AVclassifiation(out_path, PredAll1, PredAll2, VesselPredAll, DataSet=0, image_basename=''): """ predAll1: predition results of artery predAll2: predition results of vein VesselPredAll: predition results of vessel DataSet: the length of dataset image_basename: the name of saved mask """ ImgN = DataSet for ImgNumber in range(ImgN): height, width = PredAll1.shape[2:4] VesselProb = VesselPredAll[ImgNumber, 0, :, :] ArteryProb = PredAll1[ImgNumber, 0, :, :] VeinProb = PredAll2[ImgNumber, 0, :, :] VesselSeg = (VesselProb >= 0.1) & ((ArteryProb >0.2) | (VeinProb > 0.2)) # VesselSeg = (VesselProb >= 0.5) & ((ArteryProb >= 0.5) | (VeinProb >= 0.5)) crossSeg = (VesselProb >= 0.1) & ((ArteryProb >= 0.6) & (VeinProb >= 0.6)) VesselSeg = binaryPostProcessing3(VesselSeg, removeArea=100, fillArea=20) vesselPixels = np.where(VesselSeg > 0) ArteryProb2 = np.zeros((height, width)) VeinProb2 = np.zeros((height, width)) crossProb2 = np.zeros((height, width)) image_color = np.zeros((3, height, width), dtype=np.uint8) for i in range(len(vesselPixels[0])): row = vesselPixels[0][i] col = vesselPixels[1][i] probA = ArteryProb[row, col] probV = VeinProb[row, col] #probA,probV = softmax([probA,probV]) ArteryProb2[row, col] = probA VeinProb2[row, col] = probV test_use_vessel = np.zeros((height, width), np.uint8) ArteryPred2 = ((ArteryProb2 >= 0.2) & (ArteryProb2 >= VeinProb2)) VeinPred2 = ((VeinProb2 >= 0.2) & (VeinProb2 >= ArteryProb2)) ArteryPred2 = binaryPostProcessing3(ArteryPred2, removeArea=100, fillArea=20) VeinPred2 = binaryPostProcessing3(VeinPred2, removeArea=100, fillArea=20) image_color[0, :, :] = ArteryPred2 * 255 image_color[2, :, :] = VeinPred2 * 255 image_color = image_color.transpose((1, 2, 0)) #Image.fromarray(image_color).save(os.path.join(out_path, f'{image_basename[ImgNumber].split(".")[0]}_ori.png')) imgBin_vessel = ArteryPred2 + VeinPred2 imgBin_vessel[imgBin_vessel[:, :] == 2] = 1 test_use_vessel = imgBin_vessel.copy() * 255 vessel = cal_crosspoint(test_use_vessel) contours_vessel, hierarchy_c = cv2.findContours(vessel, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) # inter continuity for vessel_seg in range(len(contours_vessel)): C_vessel = np.zeros(vessel.shape, np.uint8) C_vessel = cv2.drawContours(C_vessel, contours_vessel, vessel_seg, (255, 255, 255), cv2.FILLED) cli = np.mean(VeinProb2[C_vessel == 255]) / np.mean(ArteryProb2[C_vessel == 255]) if cli < 1: image_color[ (C_vessel[:, :] == 255) & (test_use_vessel[:, :] == 255)] = [255, 0, 0] else: image_color[ (C_vessel[:, :] == 255) & (test_use_vessel[:, :] == 255)] = [0, 0, 255] loop=0 while loop<2: # out vein continuity vein = image_color[:, :, 2] contours_vein, hierarchy_b = cv2.findContours(vein, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) vein_size = [] for z in range(len(contours_vein)): vein_size.append(contours_vein[z].size) vein_size = np.sort(np.array(vein_size)) # image_color_copy = np.uint8(image_color).copy() for vein_seg in range(len(contours_vein)): judge_number = min(np.mean(vein_size),500) # cv2.putText(image_color_copy, str(vein_seg), (int(contours_vein[vein_seg][0][0][0]), int(contours_vein[vein_seg][0][0][1])), 3, 1, # color=(255, 0, 0), thickness=2) if contours_vein[vein_seg].size < judge_number: C_vein = np.zeros(vessel.shape, np.uint8) C_vein = cv2.drawContours(C_vein, contours_vein, vein_seg, (255, 255, 255), cv2.FILLED) max_diameter = np.max(Skeleton(C_vein, C_vein)[1]) image_color_copy_vein = image_color[:, :, 2].copy() image_color_copy_arter = image_color[:, :, 0].copy() # a_ori = cv2.drawContours(a_ori, contours_b, k, (0, 0, 0), cv2.FILLED) image_color_copy_vein = cv2.drawContours(image_color_copy_vein, contours_vein, vein_seg, (0, 0, 0), cv2.FILLED) # image_color[(C_cross[:, :] == 255) & (image_color[:, :, 1] == 255)] = [255, 0, 0] kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ( 4 * int(np.ceil(max_diameter)), 4 * int(np.ceil(max_diameter)))) C_vein_dilate = cv2.dilate(C_vein, kernel, iterations=1) # cv2.imwrite(path_out_3, C_vein_dilate) C_vein_dilate_judge = np.zeros(vessel.shape, np.uint8) C_vein_dilate_judge[ (C_vein_dilate[:, :] == 255) & (image_color_copy_vein == 255)] = 1 C_arter_dilate_judge = np.zeros(vessel.shape, np.uint8) C_arter_dilate_judge[ (C_vein_dilate[:, :] == 255) & (image_color_copy_arter == 255)] = 1 if (len(np.unique(C_vein_dilate_judge)) == 1) & ( len(np.unique(C_arter_dilate_judge)) != 1) & (np.mean(VeinProb2[C_vein == 255]) < 0.6): image_color[ (C_vein[:, :] == 255) & (image_color[:, :, 2] == 255)] = [255, 0, 0] # out artery continuity arter = image_color[:, :, 0] contours_arter, hierarchy_a = cv2.findContours(arter, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE) arter_size = [] for z in range(len(contours_arter)): arter_size.append(contours_arter[z].size) arter_size = np.sort(np.array(arter_size)) for arter_seg in range(len(contours_arter)): judge_number = min(np.mean(arter_size),500) if contours_arter[arter_seg].size < judge_number: C_arter = np.zeros(vessel.shape, np.uint8) C_arter = cv2.drawContours(C_arter, contours_arter, arter_seg, (255, 255, 255), cv2.FILLED) max_diameter = np.max(Skeleton(C_arter, test_use_vessel)[1]) image_color_copy_vein = image_color[:, :, 2].copy() image_color_copy_arter = image_color[:, :, 0].copy() kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ( 4 * int(np.ceil(max_diameter)), 4 * int(np.ceil(max_diameter)))) image_color_copy_arter = cv2.drawContours(image_color_copy_arter, contours_arter, arter_seg, (0, 0, 0), cv2.FILLED) C_arter_dilate = cv2.dilate(C_arter, kernel, iterations=1) # image_color[(C_cross[:, :] == 255) & (image_color[:, :, 1] == 255)] = [255, 0, 0] C_arter_dilate_judge = np.zeros(arter.shape, np.uint8) C_arter_dilate_judge[ (C_arter_dilate[:, :] == 255) & (image_color_copy_arter[:, :] == 255)] = 1 C_vein_dilate_judge = np.zeros(arter.shape, np.uint8) C_vein_dilate_judge[ (C_arter_dilate[:, :] == 255) & (image_color_copy_vein[:, :] == 255)] = 1 if (len(np.unique(C_arter_dilate_judge)) == 1) & ( len(np.unique(C_vein_dilate_judge)) != 1) & (np.mean(ArteryProb2[C_arter == 255]) < 0.6): image_color[ (C_arter[:, :] == 255) & (image_color[:, :, 0] == 255)] = [0, 0, 255] loop=loop+1 # image_basename = os.path.basename(image_basename) # Image.fromarray(image_color).save(os.path.join(out_path, f'{image_basename.split(".")[0]}.png')) # Image.fromarray(np.uint8(VesselProb*255)).save(os.path.join(out_path, f'{image_basename.split(".")[0]}_vessel.png')) return image_color