ericup/celldetection
3
1import spaces2import gradio as gr3from util import imread, imsave, copy_skimage_data4import torch5from PIL import Image, ImageDraw6import numpy as np7from os.path import join8 9 10def torch_compile(*args, **kwargs):11 def decorator(func):12 return func13 14 return decorator15 16 17torch.compile = torch_compile # temporary workaround18 19default_model = 'ginoro_CpnResNeXt101UNet-fbe875f1a3e5ce2c'20default_score_thresh = .921default_nms_thresh = np.round(np.pi / 10, 4)22default_samples = 12823default_order = 524 25examples_dir = 'examples'26copy_skimage_data(examples_dir)27examples = [28 [join(examples_dir, 'bbbc039_test_00014.png'), 'ginoro_CpnResNeXt101UNet-fbe875f1a3e5ce2c', False, default_score_thresh, False,29 default_nms_thresh, True, 64, True],30 [join(examples_dir, 'coins.png'), 'ginoro_CpnResNeXt101UNet-fbe875f1a3e5ce2c', False, default_score_thresh, False,31 default_nms_thresh, True, 64, True],32 [join(examples_dir, 'cell.png'), 'ginoro_CpnResNeXt101UNet-fbe875f1a3e5ce2c', False, default_score_thresh, False,33 default_nms_thresh, True, 64, True],34]35 36 37@spaces.GPU38def predict(39 filename, model=None,40 enable_score_threshold=False, score_threshold=.9,41 enable_nms_threshold=False, nms_threshold=0.3141592653589793,42 enable_samples=False, samples=128,43 use_label_channels=False,44 enable_order=False, order=5,45 device=None,46):47 from cpn import CpnInterface48 from prep import multi_norm49 from celldetection import label_cmap, to_h5, data, __version__50 51 global default_model52 assert isinstance(filename, str)53 54 if device is None:55 if torch.cuda.device_count():56 device = 'cuda'57 else:58 device = 'cpu'59 60 meta = dict(61 cd_version=__version__,62 filename=str(filename),63 model=model,64 device=device,65 use_label_channels=use_label_channels,66 enable_score_threshold=enable_score_threshold,67 score_threshold=float(score_threshold),68 enable_order=enable_order,69 order=order,70 enable_nms_threshold=enable_nms_threshold,71 nms_threshold=float(nms_threshold),72 )73 print(meta, flush=True)74 75 raw = img = imread(filename)76 print('Image:', img.dtype, img.shape, (img.min(), img.max()), flush=True)77 if model is None or len(str(model)) <= 0:78 model = default_model79 80 img = multi_norm(img, 'cstm-mix') # TODO81 82 kw = {}83 if enable_score_threshold:84 kw['score_thresh'] = score_threshold85 if enable_nms_threshold:86 kw['nms_thresh'] = nms_threshold87 if enable_order:88 kw['order'] = order89 if enable_samples:90 kw['samples'] = samples91 m = CpnInterface(model.strip(), device=device, **kw)92 y = m(img, reduce_labels=not use_label_channels)93 94 dst_h5 = '.'.join(filename.split('.')[:-1]) + '.h5'95 to_h5(96 dst_h5, inputs=img, **y,97 attributes=dict(inputs=meta)98 )99 100 labels = y['labels']101 vis_labels = label_cmap(labels)102 103 dst_csv = '.'.join(filename.split('.')[:-1]) + '.csv'104 data.labels2property_table(105 labels,106 "label", "area", "feret_diameter_max", "bbox", "centroid", "convex_area",107 "eccentricity", "equivalent_diameter",108 "extent", "filled_area", "major_axis_length",109 "minor_axis_length", "orientation", "perimeter",110 "solidity", "mean_intensity", "max_intensity", "min_intensity",111 intensity_image=raw112 ).to_csv(dst_csv)113 114 return vis_labels, img, dst_h5, dst_csv115 116 117with gr.Blocks(title='Cell Segmentation with Contour Proposal Networks') as app:118 with gr.Row():119 gr.Markdown("<center><strong><font size='7'>"120 "Cell Segmentation with Contour Proposal Networks 🤗</font></strong></center>")121 122 with gr.Row():123 with gr.Column():124 img = gr.components.Image(label="Upload Input Image", type="filepath", interactive=True,125 value=examples[0][0])126 with gr.Column():127 model_name = gr.components.Textbox(label='Model Name', value=default_model, max_lines=1)128 with gr.Row():129 score_thresh_ck = gr.components.Checkbox(label="Use custom Score Threshold", value=False)130 score_thresh = gr.components.Slider(minimum=0, maximum=1, label="Score Threshold",131 value=default_score_thresh)132 with gr.Row():133 nms_thresh_ck = gr.components.Checkbox(label="Use custom NMS Threshold", value=False)134 nms_thresh = gr.components.Slider(minimum=0, maximum=1, label="NMS Threshold", value=default_nms_thresh)135 # with gr.Row():136 # # The range of this would need to be model dependent137 # order_ck = gr.components.Checkbox(label="Use custom Order", value=False)138 # order = gr.components.Slider(minimum=0, maximum=1, label="Order", value=default_order)139 with gr.Row():140 samples_ck = gr.components.Checkbox(label="Use custom Sample Points", value=False)141 samples = gr.components.Slider(minimum=8, maximum=256, label="Sample Points", value=default_samples)142 with gr.Row():143 channels = gr.components.Checkbox(label="Allow overlapping objects", value=True)144 with gr.Row():145 clr = gr.Button('Reset')146 btn = gr.Button('Run')147 with gr.Row():148 with gr.Column():149 out_img = gr.Image(label="Processed Image")150 with gr.Column():151 out_vis = gr.Image(label="Label Image (random colors, transparent overlap)")152 with gr.Row():153 out_h5 = gr.File(label="Download Results as HDF5 File")154 out_csv = gr.File(label="Download Properties as CSV File")155 156 with gr.Row():157 gr.Examples(158 fn=predict,159 examples=examples,160 inputs=[img, model_name, score_thresh_ck, score_thresh, nms_thresh_ck, nms_thresh, samples_ck, samples,161 channels],162 outputs=[out_vis, out_img, out_h5, out_csv],163 cache_examples=True,164 batch=False165 )166 167 btn.click(168 predict,169 inputs=[img, model_name, score_thresh_ck, score_thresh, nms_thresh_ck, nms_thresh, samples_ck, samples,170 channels],171 outputs=[out_vis, out_img, out_h5, out_csv]172 )173 clr.click(174 lambda: (175 None, default_score_thresh, default_nms_thresh, False, False, None, None, None, False, default_samples),176 inputs=[],177 outputs=[img, score_thresh, nms_thresh, score_thresh_ck, nms_thresh_ck, out_img, out_h5, out_vis, samples_ck,178 samples]179 )180 181 with gr.Row():182 gr.Markdown("<center><font size='3'>"183 "<a href='https://github.com/FZJ-INM1-BDA/celldetection'>Visit us on GitHub</a></font></center>")184app.launch()185