| |
| |
| |
|
|
| import os |
| join = os.path.join |
| import argparse |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import tifffile as tif |
| import monai |
| from tqdm import tqdm |
| from utils.postprocess import mask_overlay |
| from monai.transforms import Activations, AddChanneld, AsChannelFirstd, AsDiscrete, Compose, EnsureTyped, EnsureType |
| from models.unicell_modules import MiT_B2_UNet_MultiHead, MiT_B3_UNet_MultiHead |
| import matplotlib.pyplot as plt |
| from skimage import io, exposure, segmentation, morphology |
| from utils.postprocess import watershed_post |
| from utils.multi_task_sliding_window_inference import multi_task_sliding_window_inference |
| import gradio as gr |
|
|
| def normalize_channel(img, lower=0.1, upper=99.9): |
| non_zero_vals = img[np.nonzero(img)] |
| percentiles = np.percentile(non_zero_vals, [lower, upper]) |
| if percentiles[1] - percentiles[0] > 0.001: |
| img_norm = exposure.rescale_intensity(img, in_range=(percentiles[0], percentiles[1]), out_range='uint8') |
| else: |
| img_norm = img |
| return img_norm |
|
|
| def preprocess(img_data): |
| if len(img_data.shape) == 2: |
| img_data = np.repeat(np.expand_dims(img_data, axis=-1), 3, axis=-1) |
| elif len(img_data.shape) == 3 and img_data.shape[-1] > 3: |
| img_data = img_data[:,:, :3] |
| else: |
| pass |
| pre_img_data = np.zeros(img_data.shape, dtype=np.uint8) |
| for i in range(3): |
| img_channel_i = img_data[:,:,i] |
| if len(img_channel_i[np.nonzero(img_channel_i)])>0: |
| pre_img_data[:,:,i] = normalize_channel(img_channel_i, lower=1, upper=99) |
| return pre_img_data |
|
|
|
|
| def inference(pre_img_data): |
| test_npy = pre_img_data/np.max(pre_img_data) |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model = MiT_B2_UNet_MultiHead(in_channels=3, out_channels=3, regress_class=1, img_size=256).to(device) |
| checkpoint = torch.load('./model.pth', map_location=torch.device(device)) |
| model.load_state_dict(checkpoint['model_state_dict']) |
| model.eval() |
| with torch.no_grad(): |
| test_tensor = torch.from_numpy(np.expand_dims(test_npy, 0)).permute(0,3,1,2).type(torch.FloatTensor).to(device) |
| |
| val_pred, val_pred_dist = multi_task_sliding_window_inference(inputs=test_tensor, roi_size=(256, 256), sw_batch_size=8, predictor=model) |
|
|
| |
| val_seg_inst = watershed_post(val_pred_dist.squeeze(1).cpu().numpy(), val_pred.squeeze(1).cpu().numpy()[:,1]) |
| test_pred_mask = val_seg_inst.squeeze().astype(np.uint16) |
| |
| |
| boundary = segmentation.find_boundaries(test_pred_mask, connectivity=1, mode='inner') |
| boundary = morphology.binary_dilation(boundary, morphology.disk(1)) |
| pre_img_data[boundary, 0] = 0 |
| pre_img_data[boundary, 1] = 255 |
| pre_img_data[boundary, 2] = 0 |
|
|
| return test_pred_mask, pre_img_data |
|
|
|
|
| def predict(img): |
| print('##########', img.name) |
| img_name = img.name |
| if img_name.endswith('.tif') or img_name.endswith('.tiff'): |
| img_data = tif.imread(img_name) |
| else: |
| img_data = io.imread(img_name) |
| if len(img_data.shape)==2: |
| pre_img_data = normalize_channel(img_data, lower=0.1, upper=99.9) |
| pre_img_data = np.repeat(np.expand_dims(pre_img_data, -1), repeats=3, axis=-1) |
| else: |
| pre_img_data = np.zeros((img_data.shape[0], img_data.shape[1], 3), dtype=np.uint8) |
| for i in range(3): |
| img_channel_i = img_data[:,:,i] |
| if len(img_channel_i[np.nonzero(img_channel_i)])>0: |
| pre_img_data[:,:,i] = normalize_channel(img_channel_i, lower=0.1, upper=99.9) |
| |
| seg_labels, seg_overlay = inference(pre_img_data) |
|
|
| tif.imwrite(join(os.getcwd(), 'segmentation.tiff'), seg_labels, compression='zlib') |
|
|
| return seg_overlay, join(os.getcwd(), 'segmentation.tiff') |
|
|
| unicell_api = gr.Interface( |
| predict, |
| inputs = gr.File(label="Input image (png, bmp, jpg, tif, tiff)"), |
| outputs = [gr.Image(label="Segmentation overlay"), gr.File(label="Download segmentation")], |
| title = "UniCell Online Demo", |
| examples=['demo.png', 'demo.tif'] |
| ) |
|
|
| unicell_api.launch() |
|
|
|
|
