Spaces:
Running on Zero
Running on Zero
| import numpy as np | |
| import gradio as gr | |
| import spaces | |
| import cv2 | |
| from cellpose import models | |
| from matplotlib.colors import hsv_to_rgb | |
| import matplotlib.pyplot as plt | |
| import os, io, base64 | |
| from PIL import Image | |
| from cellpose.io import imread, imsave | |
| import glob | |
| from huggingface_hub import hf_hub_download | |
| img = np.zeros((96, 128), dtype = np.uint8) | |
| fp0 = Image.fromarray(img) | |
| # --------------------------------------------------------------------------- | |
| # Model weights | |
| # --------------------------------------------------------------------------- | |
| MODEL_V1 = "cpsam (v1)" | |
| MODEL_V2 = "cpsam_v2" | |
| DEFAULT_MODEL = MODEL_V2 | |
| def download_weights_v1(): | |
| return hf_hub_download(repo_id="mouseland/cellpose-sam", filename="cpsam") | |
| def download_weights_v2(): | |
| # try Hugging Face first, then fall back to OSF | |
| try: | |
| return hf_hub_download(repo_id="mouseland/cellpose-sam", filename="cpsam_v2") | |
| except Exception as e: | |
| print(f"cpsam_v2 not found on Hugging Face ({e}), trying OSF") | |
| import requests | |
| fname = "cpsam_v2" | |
| url = "https://osf.io/d7c8e/download" | |
| if os.path.isfile(fname): | |
| return fname | |
| for ntries in range(10): | |
| try: | |
| r = requests.get(url, timeout=300) | |
| if r.status_code == requests.codes.ok: | |
| with open(fname, "wb") as fid: | |
| fid.write(r.content) | |
| return fname | |
| print(f"!!! Failed to download cpsam_v2 (status {r.status_code}), try {ntries+1} !!!") | |
| except Exception as e: | |
| print(f"!!! Failed to download cpsam_v2 ({e}), try {ntries+1} !!!") | |
| raise RuntimeError("Could not download cpsam_v2 weights") | |
| MODELS = {} | |
| try: | |
| MODELS[MODEL_V1] = models.CellposeModel(gpu=True, pretrained_model=download_weights_v1()) | |
| except Exception as e: | |
| print(f"Error loading {MODEL_V1}: {e}") | |
| try: | |
| MODELS[MODEL_V2] = models.CellposeModel(gpu=True, pretrained_model=download_weights_v2()) | |
| except Exception as e: | |
| print(f"Error loading {MODEL_V2}: {e}") | |
| if len(MODELS) == 0: | |
| print("No models could be loaded, exiting") | |
| exit(1) | |
| MODEL_CHOICES = list(MODELS.keys()) | |
| if DEFAULT_MODEL not in MODELS: | |
| DEFAULT_MODEL = MODEL_CHOICES[0] | |
| def get_model(model_name): | |
| if model_name not in MODELS: | |
| raise gr.Error(f"Model '{model_name}' is not available. Choose one of: {', '.join(MODEL_CHOICES)}") | |
| return MODELS[model_name] | |
| # --------------------------------------------------------------------------- | |
| # Plotting helpers | |
| # --------------------------------------------------------------------------- | |
| def plot_flows(y): | |
| Y = (np.clip(normalize99(y[0][0]),0,1) - 0.5) * 2 | |
| X = (np.clip(normalize99(y[1][0]),0,1) - 0.5) * 2 | |
| H = (np.arctan2(Y, X) + np.pi) / (2*np.pi) | |
| S = normalize99(y[0][0]**2 + y[1][0]**2) | |
| HSV = np.concatenate((H[:,:,np.newaxis], S[:,:,np.newaxis], S[:,:,np.newaxis]), axis=-1) | |
| HSV = np.clip(HSV, 0.0, 1.0) | |
| flow = (hsv_to_rgb(HSV) * 255).astype(np.uint8) | |
| return flow | |
| 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 | |
| def plot_overlay(img, masks): | |
| if img.ndim>2: | |
| img_gray = img.astype(np.float32).mean(axis=-1) | |
| else: | |
| img_gray = img.astype(np.float32) | |
| img = normalize99(img_gray) | |
| HSV = np.zeros((img.shape[0], img.shape[1], 3), np.float32) | |
| HSV[:,:,2] = np.clip(img*1.5, 0, 1.0) | |
| for n in range(int(masks.max())): | |
| ipix = (masks==n+1).nonzero() | |
| HSV[ipix[0],ipix[1],0] = np.random.rand() | |
| HSV[ipix[0],ipix[1],1] = 1.0 | |
| RGB = (hsv_to_rgb(HSV) * 255).astype(np.uint8) | |
| return RGB | |
| 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 | |
| # --------------------------------------------------------------------------- | |
| # GPU inference | |
| # --------------------------------------------------------------------------- | |
| def run_model_gpu(img, max_iter, flow_threshold, cellprob_threshold, model_name): | |
| model = get_model(model_name) | |
| masks, flows, _ = model.eval(img, niter = max_iter, flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold) | |
| return masks, flows | |
| def run_model_gpu60(img, max_iter, flow_threshold, cellprob_threshold, model_name): | |
| model = get_model(model_name) | |
| masks, flows, _ = model.eval(img, niter = max_iter, flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold) | |
| return masks, flows | |
| def run_model_gpu240(img, max_iter, flow_threshold, cellprob_threshold, model_name): | |
| model = get_model(model_name) | |
| masks, flows, _ = model.eval(img, niter = max_iter, flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold) | |
| return masks, flows | |
| import datetime | |
| from zipfile import ZipFile | |
| def cellpose_segment(filepath, resize = 1000, max_iter = 250, flow_threshold= 0.4, cellprob_threshold = 0, model_name = DEFAULT_MODEL): | |
| if not filepath: | |
| raise gr.Error("Please upload an image first.") | |
| 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) | |
| maxsize = np.max(img.shape) | |
| if maxsize<=1000: | |
| masks, flows = run_model_gpu(img, max_iter, flow_threshold, cellprob_threshold, model_name) | |
| elif maxsize < 5000: | |
| masks, flows = run_model_gpu60(img, max_iter, flow_threshold, cellprob_threshold, model_name) | |
| elif maxsize < 20000: | |
| masks, flows = run_model_gpu240(img, max_iter, flow_threshold, cellprob_threshold, model_name) | |
| else: | |
| raise ValueError("Image size must be less than 20,000") | |
| print(formatted_now, model_name, j, masks.max(), os.path.split(filepath[j])[-1]) | |
| target_size = (img_input.shape[1], img_input.shape[0]) | |
| if (target_size[0]!=img.shape[1] or target_size[1]!=img.shape[0]): | |
| # scale it back to keep the original size | |
| 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]) | |
| flows = flows[0] | |
| outpix = plot_outlines(img, masks) | |
| flows = Image.fromarray(flows) | |
| Ly, Lx = img.shape[:2] | |
| outpix = outpix.resize((Lx, Ly), resample = Image.BICUBIC) | |
| flows = flows.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, flows, b1, b2 | |
| def download_function(): | |
| b1 = gr.DownloadButton("Download masks as TIFF", visible=False) | |
| b2 = gr.DownloadButton("Download outline image as PNG", visible=False) | |
| return b1, b2 | |
| 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.newxis], [1,1,3]) | |
| elif img.ndim==3: | |
| imin = np.argmin(img.shape) | |
| if imin<2: | |
| img = np.tranpose(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]) | |
| return filepath_show, filepath, fp0, fp0 | |
| def update_button(filepath): | |
| filepath = tif_view(filepath) | |
| filepath_show = norm_path(filepath) | |
| return filepath_show, [filepath], fp0, fp0 | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| with gr.Blocks(title = "Hello", | |
| css=".gradio-container {background:purple;}") 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:white;">CellposeSAM for cellular | |
| segmentation <a style="color:#cfe7fe; font-size:14pt;" href="https://www.biorxiv.org/content/10.1101/2025.04.28.651001v1" target="_blank">[paper]</a> | |
| <a style="color:white; font-size:14pt;" href="https://github.com/MouseLand/cellpose" target="_blank">[github]</a> | |
| <a style="color:white; font-size:14pt;" href="https://www.youtube.com/watch?v=KIdYXgQemcI" target="_blank">[talk]</a> | |
| </div>""") | |
| gr.HTML("""<h4 style="color:white;">You may need to login/refresh for 5 minutes of free GPU compute per day (enough to process hundreds of images). Choose between the original 'cpsam' (v1) and 'cpsam_v2' (updated June 2026) in the model dropdown. </h4>""") | |
| input_image = gr.Image(label = "Input", type = "filepath") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| model_choice = gr.Dropdown(choices = MODEL_CHOICES, value = DEFAULT_MODEL, | |
| label = "Model", interactive = True) | |
| with gr.Row(): | |
| resize = gr.Number(label = 'max resize', value = 1000) | |
| max_iter = gr.Number(label = 'max iterations', value = 250) | |
| 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 CellposeSAM") | |
| 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 = "Outlines", type = "pil", format = 'png', value = fp0) | |
| flows = gr.Image(label = "Cellpose flows", type = "pil", format = 'png', value = fp0) | |
| sample_list = glob.glob("samples/*.png") | |
| gr.Examples(sample_list, fn = update_button, inputs=input_image, outputs = [input_image, up_btn, outlines, flows], examples_per_page=50, label = "Click on an example to try it") | |
| input_image.upload(update_button, input_image, [input_image, up_btn, outlines, flows]) | |
| up_btn.upload(update_image, up_btn, [input_image, up_btn, outlines, flows]) | |
| send_btn.click(cellpose_segment, | |
| [up_btn, resize, max_iter, flow_threshold, cellprob_threshold, model_choice], | |
| [outlines, flows, down_btn, down_btn2]) | |
| gr.HTML("""<h4 style="color:white;"> Notes:<br> | |
| <li>you can load and process 2D, multi-channel tifs. | |
| <li>the smallest dimension of a tif --> channels | |
| <li>you can upload multiple files and download a zip of the segmentations | |
| <li>install Cellpose-SAM locally for full functionality. | |
| </h4>""") | |
| demo.launch() | |