Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import gradio as gr | |
| import cv2 | |
| from cellpose import models | |
| import matplotlib.pyplot as plt | |
| import os, io | |
| from PIL import Image | |
| from cellpose.io import imread, imsave | |
| import glob | |
| import datetime | |
| from zipfile import ZipFile | |
| from huggingface_hub import hf_hub_download | |
| # --------------------------------------------------------------------------- | |
| # Download MicroAtlas model weights from Hugging Face Hub | |
| # --------------------------------------------------------------------------- | |
| # Hugging Face token: read from the HF_TOKEN secret configured in the Space. | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| def download_weights(): | |
| return hf_hub_download( | |
| repo_id="MicroAtlas/microatlas-model", | |
| filename="microatlas", | |
| token=HF_TOKEN, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Load model on CPU (runs once at startup) | |
| # --------------------------------------------------------------------------- | |
| try: | |
| fpath = download_weights() | |
| model = models.CellposeModel(gpu=False, pretrained_model=fpath) | |
| print(f"MicroAtlas model loaded from: {fpath}") | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| exit(1) | |
| # --------------------------------------------------------------------------- | |
| # Utility functions | |
| # --------------------------------------------------------------------------- | |
| def normalize99(img): | |
| X = img.copy() | |
| X = (X - np.percentile(X, 1)) / (1e-10 + np.percentile(X, 99) - np.percentile(X, 1)) | |
| return X | |
| def image_resize(img, resize=400): | |
| ny, nx = img.shape[:2] | |
| if np.array(img.shape).max() > resize: | |
| if ny > nx: | |
| nx = int(nx / ny * resize) | |
| ny = resize | |
| else: | |
| ny = int(ny / nx * resize) | |
| nx = resize | |
| shape = (nx, ny) | |
| img = cv2.resize(img, shape) | |
| img = img.astype(np.uint8) | |
| return img | |
| def plot_outlines(img, masks): | |
| img = normalize99(img) | |
| img = np.clip(img, 0, 1) | |
| outpix = [] | |
| contours, hierarchy = cv2.findContours( | |
| masks.astype(np.int32), mode=cv2.RETR_FLOODFILL, method=cv2.CHAIN_APPROX_SIMPLE | |
| ) | |
| for c in range(len(contours)): | |
| pix = contours[c].astype(int).squeeze() | |
| if len(pix) > 4: | |
| peri = cv2.arcLength(contours[c], True) | |
| approx = cv2.approxPolyDP(contours[c], 0.001, True)[:, 0, :] | |
| outpix.append(approx) | |
| figsize = (6, 6) | |
| if img.shape[0] > img.shape[1]: | |
| figsize = (6 * img.shape[1] / img.shape[0], 6) | |
| else: | |
| figsize = (6, 6 * img.shape[0] / img.shape[1]) | |
| fig = plt.figure(figsize=figsize, facecolor='k') | |
| ax = fig.add_axes([0.0, 0.0, 1, 1]) | |
| ax.set_xlim([0, img.shape[1]]) | |
| ax.set_ylim([0, img.shape[0]]) | |
| ax.imshow(img[::-1], origin='upper', aspect='auto') | |
| if outpix is not None: | |
| for o in outpix: | |
| ax.plot(o[:, 0], img.shape[0] - o[:, 1], color=[1, 0, 0], lw=1) | |
| ax.axis('off') | |
| buf = io.BytesIO() | |
| fig.savefig(buf, bbox_inches='tight') | |
| buf.seek(0) | |
| pil_img = Image.open(buf) | |
| plt.close(fig) | |
| return pil_img | |
| # --------------------------------------------------------------------------- | |
| # CPU inference | |
| # --------------------------------------------------------------------------- | |
| def run_model(img, flow_threshold, cellprob_threshold): | |
| masks, flows, _ = model.eval( | |
| img, | |
| channels=None, | |
| diameter=None, | |
| bsize=256, | |
| flow_threshold=flow_threshold, | |
| cellprob_threshold=cellprob_threshold, | |
| ) | |
| return masks, flows | |
| # --------------------------------------------------------------------------- | |
| # Main segmentation pipeline | |
| # --------------------------------------------------------------------------- | |
| def microatlas_segment(filepath, resize=512, flow_threshold=0.4, cellprob_threshold=0): | |
| zip_path = os.path.splitext(filepath[-1])[0] + "_masks.zip" | |
| with ZipFile(zip_path, 'w') as myzip: | |
| for j in range(len(filepath)): | |
| now = datetime.datetime.now() | |
| formatted_now = now.strftime("%Y-%m-%d %H:%M:%S") | |
| img_input = imread(filepath[j]) | |
| img = image_resize(img_input, resize=resize) | |
| masks, flows = run_model(img, flow_threshold, cellprob_threshold) | |
| print(formatted_now, j, masks.max(), os.path.split(filepath[j])[-1]) | |
| # Scale masks back to original size | |
| target_size = (img_input.shape[1], img_input.shape[0]) | |
| if target_size[0] != img.shape[1] or target_size[1] != img.shape[0]: | |
| masks_rsz = cv2.resize( | |
| masks.astype('uint16'), target_size, interpolation=cv2.INTER_NEAREST | |
| ).astype('uint16') | |
| else: | |
| masks_rsz = masks.copy() | |
| fname_masks = os.path.splitext(filepath[j])[0] + "_masks.tif" | |
| imsave(fname_masks, masks_rsz) | |
| myzip.write(fname_masks, arcname=os.path.split(fname_masks)[-1]) | |
| # Generate visualization (based on last image) | |
| outpix = plot_outlines(img, masks) | |
| Ly, Lx = img.shape[:2] | |
| outpix = outpix.resize((Lx, Ly), resample=Image.BICUBIC) | |
| fname_out = os.path.splitext(filepath[-1])[0] + "_outlines.png" | |
| outpix.save(fname_out) | |
| if len(filepath) > 1: | |
| b1 = gr.DownloadButton(visible=True, value=zip_path) | |
| else: | |
| b1 = gr.DownloadButton(visible=True, value=fname_masks) | |
| b2 = gr.DownloadButton(visible=True, value=fname_out) | |
| return outpix, b1, b2 | |
| # --------------------------------------------------------------------------- | |
| # UI helpers | |
| # --------------------------------------------------------------------------- | |
| def tif_view(filepath): | |
| fpath, fext = os.path.splitext(filepath) | |
| if fext in ['.tiff', '.tif']: | |
| img = imread(filepath[-1]) | |
| if img.ndim == 2: | |
| img = np.tile(img[:, :, np.newaxis], [1, 1, 3]) | |
| elif img.ndim == 3: | |
| imin = np.argmin(img.shape) | |
| if imin < 2: | |
| img = np.transpose(img, [2, imin]) | |
| else: | |
| raise ValueError("TIF cannot have more than three dimensions") | |
| Ly, Lx, nchan = img.shape | |
| imgi = np.zeros((Ly, Lx, 3)) | |
| nn = np.minimum(3, img.shape[-1]) | |
| imgi[:, :, :nn] = img[:, :, :nn] | |
| imsave(filepath, imgi) | |
| return filepath | |
| def norm_path(filepath): | |
| img = imread(filepath) | |
| img = normalize99(img) | |
| img = np.clip(img, 0, 1) | |
| fpath, fext = os.path.splitext(filepath) | |
| filepath = fpath + '.png' | |
| pil_image = Image.fromarray((255. * img).astype(np.uint8)) | |
| pil_image.save(filepath) | |
| return filepath | |
| def update_image(filepath): | |
| for f in filepath: | |
| f = tif_view(f) | |
| filepath_show = norm_path(filepath[-1]) | |
| fp0 = Image.fromarray(np.zeros((96, 128), dtype=np.uint8)) | |
| return filepath_show, filepath, fp0 | |
| def update_button(filepath): | |
| filepath = tif_view(filepath) | |
| filepath_show = norm_path(filepath) | |
| fp0 = Image.fromarray(np.zeros((96, 128), dtype=np.uint8)) | |
| return filepath_show, [filepath], fp0 | |
| # --------------------------------------------------------------------------- | |
| # Gradio UI | |
| # --------------------------------------------------------------------------- | |
| fp0 = Image.fromarray(np.zeros((96, 128), dtype=np.uint8)) | |
| with gr.Blocks( | |
| title="MicroAtlas Cell Segmentation", | |
| ) as demo: | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.HTML(""" | |
| <div style="font-family:'Times New Roman', 'Serif'; font-size:20pt; font-weight:bold; text-align:center; color:#333;"> | |
| Large‑Scale Unlabeled Microscopy Images Empower a Generalizable Cell Segmentation Foundation Model for Versatile Biological Analysis<br> | |
| <a style="color:#0066cc; font-size:14pt;" href="https://huggingface.co/datasets/MicroAtlas/MicroAtlas-2B" target="_blank">[dataset]</a> | |
| <a style="color:#333; font-size:14pt;" href="https://github.com/Luffy03/MicroAtlas" target="_blank">[github]</a> | |
| </div>""") | |
| gr.HTML("""<h4 style="color:#333;"> | |
| MicroAtlas is a foundation model trained on large-scale unlabeled microscopy images. | |
| </h4>""") | |
| gr.HTML("""<h4 style="color:#333;"> | |
| For fast GPU inference, clone the code and run locally: | |
| <a style="color:#0066cc;" href="https://github.com/Luffy03/MicroAtlas" target="_blank">github.com/Luffy03/MicroAtlas</a> | |
| </h4>""") | |
| input_image = gr.Image(label="Input", type="filepath") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| with gr.Row(): | |
| resize = gr.Number(label='max resize', value=512) | |
| flow_threshold = gr.Number(label='flow threshold', value=0.4) | |
| cellprob_threshold = gr.Number(label='cellprob threshold', value=0) | |
| up_btn = gr.UploadButton( | |
| "Multi-file upload (png, jpg, tif etc)", | |
| visible=True, | |
| file_count="multiple" | |
| ) | |
| with gr.Column(scale=1): | |
| send_btn = gr.Button("Run MicroAtlas") | |
| down_btn = gr.DownloadButton("Download masks (TIF)", visible=False) | |
| down_btn2 = gr.DownloadButton("Download outlines (PNG)", visible=False) | |
| with gr.Column(scale=2): | |
| outlines = gr.Image(label="Segmentation", type="pil", format='png', value=fp0) | |
| # Example images (if you have a samples/ folder in the Space) | |
| sample_list = glob.glob("samples/*.png") | |
| if sample_list: | |
| gr.Examples( | |
| sample_list, | |
| fn=update_button, | |
| inputs=input_image, | |
| outputs=[input_image, up_btn, outlines], | |
| examples_per_page=50, | |
| label="Click on an example to try it" | |
| ) | |
| input_image.upload(update_button, input_image, [input_image, up_btn, outlines]) | |
| up_btn.upload(update_image, up_btn, [input_image, up_btn, outlines]) | |
| send_btn.click( | |
| microatlas_segment, | |
| [up_btn, resize, flow_threshold, cellprob_threshold], | |
| [outlines, down_btn, down_btn2] | |
| ) | |
| gr.HTML("""<h4 style="color:#333;"> Notes:<br> | |
| <li>This Space runs on <b>CPU</b> — expect ~1–2 minutes per image. | |
| <li>For production use, clone <a style="color:#0066cc;" href="https://github.com/Luffy03/MicroAtlas" target="_blank">github.com/Luffy03/MicroAtlas</a> and run locally with GPU. | |
| <li>You can load and process 2D, multi-channel tifs. | |
| <li>You can upload multiple files and download a zip of the segmentations. | |
| </h4>""") | |
| demo.launch(css=".gradio-container {background: white;}", ssr_mode=False) | |