Merge branch 'main' of https://huggingface.co/spaces/mouseland/cellpose

app.py

@@ -9,8 +9,19 @@ import os, io, base64
 from PIL import Image 
 from cellpose.io import imread, imsave
 
+from huggingface_hub import hf_hub_download
+
+img = np.zeros((96, 128))
+fp0 = "0.png"
+imsave(fp0, img)
+
 # @title Data retrieval
-def download_weights():
+def download_weights():    
+    return hf_hub_download(repo_id="mouseland/cellpose-sam", filename="cpsam")
+    
+    #os.system("wget -q https://huggingface.co/mouseland/cellpose-sam/resolve/main/cpsam")
+
+def download_weights_old():
     import os, requests
     
     fname = ['cpsam']
@@ -19,20 +30,24 @@ def download_weights():
     
     for j in range(len(url)):
       if not os.path.isfile(fname[j]):
-        try:
-          r = requests.get(url[j])
-        except requests.ConnectionError:
-          print("!!! Failed to download data !!!")
-        else:
-          if r.status_code != requests.codes.ok:
-            print("!!! Failed to download data !!!")
-          else:
-            with open(fname[j], "wb") as fid:
-              fid.write(r.content)
+        ntries = 0
+        while ntries<10:
+            try:
+              r = requests.get(url[j])
+            except:
+                print("!!! Failed to download data !!!")
+                ntries += 1 
+                print(ntries)
+            
+      if r.status_code != requests.codes.ok:
+        print("!!! Failed to download data !!!")
+      else:
+        with open(fname[j], "wb") as fid:
+          fid.write(r.content)
 
 try:
-    download_weights()
-    model = models.CellposeModel(gpu=True, pretrained_model="cpsam")
+    fpath = download_weights()
+    model = models.CellposeModel(gpu=True, pretrained_model = fpath)
 except Exception as e:
     print(f"Error loading model: {e}")
     exit(1)
@@ -51,6 +66,7 @@ def plot_flows(y):
     return flow
 
 def plot_outlines(img, masks):
+    img = normalize99(img)
     outpix = []
     contours, hierarchy = cv2.findContours(masks.astype(np.int32), mode=cv2.RETR_FLOODFILL, method=cv2.CHAIN_APPROX_SIMPLE)
     for c in range(len(contours)):
@@ -94,7 +110,12 @@ def plot_outlines(img, masks):
     return pil_img
 
 def plot_overlay(img, masks):
-    img = normalize99(img.astype(np.float32).mean(axis=-1))
+    if img.ndim>2:
+        img_gray = img.astype(np.float32).mean(axis=-1)
+    else:
+        img_gray = img.astype(np.float32)
+        
+    img = normalize99(img_gray)
     img -= img.min()
     img /= img.max()
     HSV = np.zeros((img.shape[0], img.shape[1], 3), np.float32)
@@ -146,66 +167,76 @@ def run_model_gpu1000(img):
     masks, flows, _ = model.eval(img)#, channels = [0,0])
     return masks, flows
 
-#@spaces.GPU(duration=10)
-def cellpose_segment(img_pil, resize = 400):
-    img_input = imread(img_pil)
-    #img_input = np.array(img_pil)
-    img = image_resize(img_input, resize = resize)
-
-    resize = np.max(img.shape)
-    if resize<1000:
-        masks, flows = run_model_gpu(img)
-    elif resize < 5000:
-        masks, flows = run_model_gpu60(img)
-    elif resize < 20000:
-        masks, flows = run_model_gpu240(img)
-    else:
-        raise ValueError("Image size must be less than 20,000")
-    
+
+from zipfile import ZipFile
+def cellpose_segment(filepath, resize = 1000):
+
+    zip_path = os.path.splitext(filepath[-1])[0]+"_masks.zip"
+    #zip_path = 'masks.zip'
+    with ZipFile(zip_path, 'w') as myzip:
+        for j in range((len(filepath))):
+            print(j)
+            img_input = imread(filepath[j])
+            #img_input = np.array(img_pil)
+            img = image_resize(img_input, resize = resize)
+            
+            maxsize = np.max(img.shape)
+            if maxsize<1000:
+                masks, flows = run_model_gpu(img)
+            elif maxsize < 5000:
+                masks, flows = run_model_gpu60(img)
+            elif maxsize < 20000:
+                masks, flows = run_model_gpu240(img)
+            else:
+                raise ValueError("Image size must be less than 20,000")
         
+            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 orignal 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])
+            
+    
     #masks, flows, _ = model.eval(img, channels=[0,0])
     flows = flows[0]
     # masks = np.zeros(img.shape[:2])
     # flows = np.zeros_like(img)
 
     outpix = plot_outlines(img, masks)
-    overlay = plot_overlay(img, masks)
-    
-    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 orignal size
-        masks = cv2.resize(masks.astype('uint16'), target_size, interpolation=cv2.INTER_NEAREST).astype('uint16')
-        #flows = cv2.resize(flows.astype('float32'), target_size).astype('uint8')
+    #overlay = plot_overlay(img, masks)
     
+        
     
     #crand = .2 + .8 * np.random.rand(np.max(masks.flatten()).astype('int')+1,).astype('float32')
     #crand[0] = 0
 
-    overlay = Image.fromarray(overlay)
+    #overlay = Image.fromarray(overlay)
     flows = Image.fromarray(flows)
 
     Ly, Lx = img.shape[:2]
-    c = Lx
     outpix = outpix.resize((Lx, Ly), resample  = Image.BICUBIC)
-    overlay = overlay.resize((Lx, Ly), resample  = Image.BICUBIC)
+    #overlay = overlay.resize((Lx, Ly), resample  = Image.BICUBIC)
     flows = flows.resize((Lx, Ly), resample  = Image.BICUBIC)
 
-    #masks = Image.fromarray(255. * crand[masks])
-    #pil_masks = Image.fromarray(masks.astype('int32'))
-    #pil_masks.save(fname_mask)
-
-    fname_out  = os.path.splitext(img_pil)[0]+"_outlines.png"
-    fname_masks = os.path.splitext(img_pil)[0]+"_masks.tif"
-
-    imsave(fname_masks, masks)
-    
-    
+    fname_out  = os.path.splitext(filepath[-1])[0]+"_outlines.png"
     outpix.save(fname_out) #"outlines.png")
 
-    b1 = gr.DownloadButton(visible=True, value = fname_masks)
+    fname_flows  = os.path.splitext(filepath[-1])[0]+"_flows.png"
+    flows.save(fname_flows) #"outlines.png")
+
+    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) #"outlines.png")
     
-    return outpix, overlay, flows, b1, b2
+    return fname_out, fname_flows, b1, b2
 
 # Gradio Interface
 #iface = gr.Interface(
@@ -221,41 +252,103 @@ def download_function():
     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]
+        
+        #filepath = fpath+'.png'
+        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)
+    #imsave(filepath, pil_image)
+    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
+    
 with gr.Blocks(title = "Hello", 
                css=".gradio-container {background:purple;}") as demo:
 
+    #filepath = ""
     with gr.Row():
         with gr.Column(scale=2):
-            gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:white;">Cellpose-SAM for cellular segmentation</div>""")
-            gr.HTML("""<h4 style="color:white;">You may need to refresh/login for 5 minutes of free GPU compute/day. </h4>""")
-            gr.HTML("""<h4 style="color:white;">"pip install cellpose" for full functionality. </h4>""")
-
-            input_image = gr.Image(label = "Input image", type = "filepath")
+            gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:20pt; font-weight:bold; text-align:center; color:white;">Cellpose-SAM for cellular 
+            segmentation <a style="color:#cfe7fe; font-size:14pt;" href="https://github.com/MouseLand/cellpose" target="_blank">[paper]</a> 
+            <a style="color:white; font-size:14pt;" href="https://github.com/MouseLand/cellpose" target="_blank">[github]</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). </h4>""")
+            
+            input_image = gr.Image(label = "Input", type = "filepath")
 
             with gr.Row():
-                resize = gr.Number(label = 'max resize', value = 400)
-                send_btn = gr.Button("Run Cellpose-SAM")
-                
-            with gr.Row():
-                down_btn = gr.DownloadButton("Download masks (TIF)", visible=False)            
-                down_btn2 = gr.DownloadButton("Download outlines (PNG)", visible=False)
-
-            gr.HTML("""<a style="color:white;" href="https://github.com/MouseLand/cellpose" target="_blank">github page for cellpose</a>""")
-            gr.HTML("""<a style="color:white;" href="https://github.com/MouseLand/cellpose" target="_blank">Cellpose-SAM paper</a>""")
-
+                with gr.Column(scale=1):                    
+                    with gr.Row():
+                        resize = gr.Number(label = 'max resize', value = 1000)
+                        up_btn = gr.UploadButton("Multi-file upload (png, jpg, tif etc)", visible=True, file_count = "multiple")                        
+                    
+                    #gr.HTML("""<h4 style="color:white;"> Note2: Only the first image of a tif will display the segmentations, but you can download segmentations for all planes. </h4>""")
+                    
+                with gr.Column(scale=1):
+                    send_btn = gr.Button("Run Cellpose-SAM")
+                    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 = "filepath", format = 'png', value = fp0) #, width = "50vw", height = "20vw")
+            #img_overlay = gr.Image(label = "Overlay", type = "pil", format = 'png') #, width = "50vw", height = "20vw")
+            flows = gr.Image(label = "Cellpose flows", type = "filepath", format = 'png', value = fp0) #, width = "50vw", height = "20vw")
 
-        with gr.Column(scale=2):            
-            img_outlines = gr.Image(label = "Outlines", type = "pil", format = 'png') #, width = "50vw", height = "20vw")
-            img_overlay = gr.Image(label = "Overlay", type = "pil", format = 'png') #, width = "50vw", height = "20vw")
-            flows = gr.Image(label = "Cellpose flows", type = "pil", format = 'png') #, width = "50vw", height = "20vw")
-            #masks = gr.Image(label = "Output image", type = "numpy")
-    
+            
+    sample_list = []
+    for j in range(23):
+        sample_list.append("samples/img%0.2d.png"%j)
+        
+    gr.Examples(sample_list, fn = update_button, inputs=input_image, outputs = [input_image, up_btn, outlines, flows], examples_per_page=25, 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(fn=cellpose_segment, inputs=[input_image, resize], outputs=[img_outlines, img_overlay, flows, down_btn, down_btn2])
+    send_btn.click(cellpose_segment, [up_btn, resize], [outlines, flows, down_btn, down_btn2])
 
     #down_btn.click(download_function, None, [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 load multiple files and download a zip of the segmentations
+                    <li>install Cellpose-SAM locally for full functionality.
+                    </h4>""")
     
-    
-
+    # <li>the smallest dimension of a tif --> channels 
+    # <li>you can load multiple files and download a zip of the segmentations  
+                    
 demo.launch()

requirements.txt

@@ -1,3 +1,3 @@
 gradio
-git+https://github.com/mouseland/cellpose_dev
+git+https://github.com/MouseLand/cellpose_dev
 matplotlib