import sys import os import glob import cv2 import imageio import xmlsettings import numpy as np import pandas as pd from pyntcloud import PyntCloud from depth_image import DepthImage from general_utils import make_dir from skimage.color import label2rgb from cvtools import img_fill, enlarge_bbox, auto_canny from skimage.measure import label, regionprops DEBUG = False NOISE_FLOOR = 10 def roi_detector(prefix, output_path, im_in, depth_img, min_area=300, size_classes=[0.1, 0.2, 0.4, 0.8, 1.6, 3.2, 6.4, 12.8], save_rois=False, cfg=None, ): # scale depth mask to image size depth_mask = depth_img.get_depth_mask() z_img = depth_img.depth_images[2] if not im_in.shape == depth_mask.shape: depth_mask = cv2.resize(depth_mask, im_in.shape) z_img = cv2.resize(z_img, im_in.shape) # fill in missing parts of depth mask from img_in kernel = np.ones((cfg.get("DepthMaskKernel", 21), cfg.get("DepthMaskKernel", 21)), np.uint8) depth_mask = cv2.morphologyEx(depth_mask, cv2.MORPH_CLOSE, kernel) depth_mask = img_fill(depth_mask.astype('uint8') * 255) / 255 # label connected components in the depth image label_image = label(depth_mask, connectivity=2) if cfg.get("Debug", "False").lower() == "true": cv2.imshow('label_image', label2rgb(label_image)) cv2.waitKey() rois = [] particle_stats = np.zeros((1, len(size_classes))) depths = [] im_out = cv2.cvtColor(im_in.copy(), cv2.COLOR_GRAY2RGB) counter = 0 valid_roi_counter = 0 esds = [] all_regions = regionprops(label_image, im_in) for region in all_regions: coords = region['coords'] if cfg.get("EnlargeBox", 0) > 0: box = enlarge_bbox(region['bbox'], cfg.get("EnlargeBox", 0), im_in) else: box = region['bbox'] depth_val = np.mean(z_img[tuple(coords[:, 0]), tuple(coords[:, 1])]) eff_res = cfg.get("PixelSize", 101.0) * (cfg.get("FocusRange", 75.0) - depth_val) / cfg.get("FocusRange", 75.0) # Put into size bin using ESD esd = 2*np.sqrt(region['area']/np.pi) * eff_res / 1000 # ESD in mm esds.append(esd) # print(esd) if region['area'] > min_area: roi = {} #roi['image'] = region['intensity_image'] roi['image'] = im_in[box[0]:box[2], box[1]:box[3]] roi['file_prefix'] = prefix min_image = np.min(roi['image'][roi['image'] > 0]) roi['image'] = 255*(roi['image'].astype('float') - min_image) / np.max(roi['image']) roi['image'][roi['image'] < 0] = 0 roi['image'] = roi['image'].astype('uint8') im_out = cv2.rectangle(im_out, (box[1], box[0]), (box[3], box[2]), (0, 255, 0), 2 ) valid_roi_counter += 1 if save_rois: imageio.imwrite( os.path.join(output_path, prefix + '-' + str(counter) + '-' + str(box[1]) + '-' + str(box[0]) + '-' + str(box[3]) + '-' + str(box[2]) + '-' + str(int(depth_val)) + '.png' ), roi['image'] ) counter += 1 h, n = np.histogram(esds, bins=size_classes) return im_out, valid_roi_counter, h, len(all_regions) def preproc_image_pair(file_prefix, depth_range=[0, 255], cfg=None): # load the total focus image tf_image = imageio.imread(file_prefix+'-TotalFocus.png') tf_image = np.flipud(tf_image) # Step 0 : generate depth mask from 3D image using one of the channels depth_img = DepthImage(file_prefix) mask = depth_img.get_depth_mask() # step 1 : Resize the mask to match the dims of the total focus image tf_width = tf_image.shape[1] tf_height = tf_image.shape[0] mask = cv2.resize(mask, (tf_width, tf_height)) # step 3 : convert image to 8-bits tf_image = (tf_image/256).astype('uint8') # step 2 : mask out artifacts in the total focus image # the assumption here is that if there is no depth calculated for that pixel # it should be excluded from the image. masked_tf = mask * tf_image masked_tf = (masked_tf).astype('uint8') # cv2.imshow('edge', cv2.resize(mask, (1024, 1024))) # cv2.waitKey() # step 4: mask out the background artifacts around valid depths masked_tf[masked_tf < NOISE_FLOOR] = NOISE_FLOOR tf_image[tf_image < NOISE_FLOOR] = NOISE_FLOOR #lum_mask = masked_tf.copy() #lum_mask[lum_mask > 0] = 1 # These last two step may not be needed # step 5 : segment the masked total focus image and fill holes in the segmentation edge_image = auto_canny(masked_tf, sigma=cfg.get("AutoCannySigma", 0.33)) kernel = np.ones((9, 9), np.uint8) edge_image = cv2.morphologyEx(edge_image, cv2.MORPH_CLOSE, kernel) edge_filled = img_fill(edge_image).astype('uint8') edge_filled[edge_filled > 0] = 1 # step 6: mask the depth images with the segmented total focus image depth_img.mask(edge_filled) if cfg.get("Debug", "False").lower() == "true": cv2.imshow('depth', cv2.resize(depth_img.depth_images[3], (1024, 1024))) cv2.waitKey() return tf_image, masked_tf, depth_img if __name__=="__main__": if len(sys.argv) < 5: print('Usage: python RaytrixParticleProc.py [settings_file] [input dir] [pattern] [output_dir]') exit() focus_images = sorted(glob.glob(os.path.join(sys.argv[2], sys.argv[3] + '-TotalFocus.png'))) settings_path = sys.argv[1] output_dir = sys.argv[4] # load the config file cfg = xmlsettings.XMLSettings(settings_path) if len(focus_images) > 0: # create output dirs roi_dir = make_dir(output_dir, 'rois') stats_dir = make_dir(output_dir, 'stats') images_dir = make_dir(output_dir, 'images') clouds_dir = make_dir(output_dir, 'clouds') stats_file = os.path.join(stats_dir, 'basic_stats.csv') total_rois = 0 for img in focus_images: sub_fmt = "{:010d}".format(int(total_rois / 5000)*5000) roi_sub_dir = os.path.join(roi_dir, sub_fmt) if not os.path.exists(roi_sub_dir): make_dir(roi_dir, sub_fmt) prefix = img[0:-15] # remove the '-TotalFocus.png' print(prefix) tf_image, masked_tf, depth_img = preproc_image_pair(prefix, cfg=cfg) cloud = depth_img.generate_point_cloud() if cfg.get("Debug", "False").lower() == "true": cv2.imshow('masked_tf', cv2.resize(masked_tf, (1024, 1024))) cv2.imshow('masked depth', depth_img.depth_images[3]) cv2.waitKey() else: im_out, n_rois, h, n_particles = \ roi_detector(os.path.basename(prefix), roi_sub_dir, tf_image, depth_img, cfg=cfg, save_rois=True) total_rois += n_rois with open(stats_file, "a+") as f: output = os.path.basename(prefix) + "," output += str(n_particles) + "," for i in range(len(h)): output += str(h[i]) + "," output += str(depth_img.num_particles) + "," output += str(depth_img.nnstat) f.write(output + "\n") imageio.imwrite(os.path.join(images_dir, os.path.basename(prefix) + '-masked-tf.png'), im_out) if np.prod(cloud.shape) != 0: cloud = PyntCloud(pd.DataFrame( # same arguments that you are passing to visualize_pcl data=cloud, columns=["x", "y", "z", "size"]) ) cloud.to_file(os.path.join(clouds_dir, os.path.basename(prefix) + ".ply"))