Upload folder using huggingface_hub

README.md

@@ -1,12 +1,6 @@
 ---
-title: Uw Ct Seg
-emoji: 👁
-colorFrom: blue
-colorTo: blue
+title: uw-ct-seg
+app_file: ct_seg.py
 sdk: gradio
 sdk_version: 4.44.0
-app_file: app.py
-pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

ct_seg.py

@@ -0,0 +1,316 @@
+# --------------------------------------------------------
+# BiomedSeg
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Yu Gu (yugu1@microsoft.com), Theo Zhao (theodorezhao@microsoft.com)
+# --------------------------------------------------------
+
+import os
+import sys
+
+this_file_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.join(this_file_dir, "../ct_seg"))
+import json
+import warnings
+import PIL
+from PIL import Image
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import monai
+import cv2
+import math
+import gradio as gr
+import torch
+import argparse
+import imageio
+import numpy as np
+import scipy
+
+from torchvision import transforms
+from models import dinov2_vitl_transunet
+from class_dict import class_dict, dataset_class
+from transforms import _MEAN, _STD
+from monai import transforms as monai_transforms
+from scipy.ndimage import label
+
+id2label = {v: k for k, v in class_dict.items()}
+np.random.seed(0)
+id2color = {k: list(np.random.choice(range(256), size=3)) for k,v in id2label.items()}
+
+
+def clean_mask(X):
+    """
+    Cleans the mask for labels 1 and 2 by keeping only the largest connected component for each label.
+    
+    Parameters:
+    X (numpy.ndarray): Volumetric mask of shape [N, 1, W, H] with values 0 (background), 1, or 2.
+    
+    Returns:
+    numpy.ndarray: Cleaned volumetric mask with the same shape as X.
+    """
+    # Extract the volume data (assuming N is the depth dimension)
+    if X.ndim == 4:
+        volume = X[:, 0, :, :]  # Shape: [N, W, H]
+    else:
+        volume = X
+    
+    for label_value in [1, 2, 10]:
+        # Create a binary mask for the current label
+        mask = (volume == label_value)
+        if not np.any(mask):
+            continue  # Skip if the label is not present
+        
+        # Define connectivity for 3D connected components
+        structure = np.ones((3, 3, 3), dtype=int)
+        
+        # Label connected components
+        labeled_mask, num_features = label(mask, structure=structure)
+        if num_features == 0:
+            continue  # No connected components found
+        
+        # Compute sizes of all connected components
+        component_sizes = np.bincount(labeled_mask.ravel())
+        component_sizes[0] = 0  # Ignore the background
+        
+        # Find the label of the largest connected component
+        largest_component_label = component_sizes.argmax()
+        
+        # Create a mask for the largest connected component
+        largest_component_mask = (labeled_mask == largest_component_label)
+        
+        # Remove all other components of the current label
+        volume[mask] = 0  # Set all pixels of the current label to background
+        volume[largest_component_mask] = label_value  # Restore the largest component
+    
+    # Update the original mask
+    if X.ndim == 4:
+        X[:, 0, :, :] = volume
+    else:
+        X = volume
+    return X
+
+
+def parse_option():
+    parser = argparse.ArgumentParser('SEEM Demo', add_help=False)
+    parser.add_argument('--model_path', default="ckpt/model_19.pth", metavar="FILE", help='path to model file')
+    # parser.add_argument('--model_path', default="ckpt/uw_seg_heart.pth", metavar="FILE", help='path to model file')
+    cfg = parser.parse_args()
+    return cfg
+
+'''
+build args
+'''
+cfg = parse_option()
+
+pretrained_pth = cfg.model_path
+
+def load_tif_images(file_path):
+    vol = imageio.imread(file_path)
+    if np.max(vol) <= 1:
+        vol = vol * 255
+    return vol
+
+def overlay_image_with_mask(image, segmentation_map, path='test.png', ax=None):
+    color_seg = np.zeros((segmentation_map.shape[0], segmentation_map.shape[1], 3), dtype=np.uint8) # height, width, 3
+    for label, color in id2color.items():
+        color_seg[segmentation_map == label, :] = color
+
+    # Show image + mask
+    img = np.array(image) * 0.5 + color_seg * 0.5
+    img = img.astype(np.uint8)
+    return img
+
+def resize_volume(vol, size, max_frames, nearest_neighbor=False):
+    W, H, F = vol.shape
+
+    zoom_rate = size / W
+    vol_reshape = scipy.ndimage.zoom(
+        vol, (zoom_rate, zoom_rate, zoom_rate), order=3 if not nearest_neighbor else 0
+    )
+    resizeW, resizeH, resizeF = vol_reshape.shape
+    if resizeF > max_frames:
+        vol_reshape = vol_reshape[:, :, :max_frames]
+        resizeF = max_frames
+    else:
+        resized_max_fr = int(math.ceil(max_frames * zoom_rate))
+        vol_reshape = np.concatenate([vol_reshape, np.zeros((resizeW, resizeH, resized_max_fr - resizeF))], axis=-1)
+    return vol_reshape, resizeF, zoom_rate
+
+val_transform = monai_transforms.Compose([monai_transforms.Resized(keys=['image'], spatial_size=(256, 256), mode=['bilinear'])])
+def process_volume(vol: np.ndarray, keep_frames: Callable=lambda x: x > 0.025):
+    initial_resize = monai.transforms.ResizeWithPadOrCrop((512, 512))
+    transform = monai.transforms.CropForeground(keys=["pixel_values"], source_key="pixel_values", return_coords=True)
+    crop_vol, start_coords, end_coords = transform(vol)
+    keep_frames = np.where(keep_frames(np.mean(np.mean(crop_vol, axis=-1), axis=-1)))[0]
+    crop_vol = crop_vol[keep_frames]
+    W, H, F = crop_vol.shape
+    proc_vol = cv2.equalizeHist(crop_vol.reshape(W, -1).astype(np.uint8)).reshape(W, H, F)
+    proc_vol = initial_resize(proc_vol).detach().cpu().numpy().transpose((1, 2, 0))
+    proc_vol, max_fr = resize_volume(proc_vol, 256, max_frames=512)[:2]
+
+    images = []
+    for i in range(proc_vol.shape[2]):
+        image = torch.from_numpy(proc_vol[:, :, i]).unsqueeze(0)
+        image_transformed = val_transform({"image": image})["image"]
+        images.append(image_transformed)
+    images = torch.stack(images)
+    if images.max() > 1:
+        images = images / 255.0
+    # make the images three channels
+    images = images.repeat(1, 3, 1, 1)
+    for c in range(len(_MEAN)):
+        images[:, c, :, :] = (images[:, c, :, :] - _MEAN[c]) / _STD[c]
+    return images, max_fr
+
+def untransform(img):
+    for c in range(len(_MEAN)):
+        img[c] = img[c] * _STD[c] + _MEAN[c]
+    if img.max() <= 1:
+        img = img * 255
+    return img.long()
+
+def process_ct(ct_path: str):
+    vol = load_tif_images(ct_path)
+    images, frame_indices = process_volume(vol, keep_frames=lambda x: x > 0.025)
+    return images, frame_indices
+
+# Ensure the example file is in the same directory or provide a relative path
+examples = [["demo/CTseg_57_raw.tif"],
+            ["demo/CTrec-don_1101.tif"]]
+
+'''
+build model
+'''
+class_names = dataset_class["uwseg"]
+class_ids = [class_dict[class_name] for class_name in class_names]
+model = dinov2_vitl_transunet(pretrained="", num_classes=len(class_dict), img_size=256)
+state_dict = torch.load(pretrained_pth)
+model.load_state_dict(state_dict)
+model = model.cuda()
+
+@torch.no_grad()
+def inference(image_input):
+    if isinstance(image_input, str):
+        # image_input is a file path
+        file_path = image_input
+    else:
+        # image_input is a gr.File object
+        file_path = image_input.name
+    images, frame_indices = process_ct(file_path)
+    with torch.no_grad():
+        with torch.cuda.amp.autocast(dtype=torch.float16):
+            logits = model(images.cuda())
+    for j in range(len(class_dict)):
+        if j + 1 not in class_ids:
+            logits[:, j] = -1000
+    pred = torch.argmax(logits, dim=1) + 1
+    pred_mask = (torch.max(logits, dim=1)[0] > 0)
+    pred = pred_mask * pred
+    pred[frame_indices:] = 0
+    pred = torch.from_numpy(clean_mask(pred.cpu().numpy()))
+    volume_size = torch.sum(pred==2).item()
+    # 1 pixel = 1 mm^2, change to cm^3
+    volume_size = volume_size / 1000
+
+    # Compute the size of the segmented mask for each slice
+    sizes = pred.view(pred.shape[0], -1).sum(dim=1).cpu().numpy()
+
+    segmentation_results = []
+    raw_images = []
+    for i in range(len(images)):
+        images[i] = untransform(images[i])
+        raw_image = Image.fromarray(images[i].cpu().permute(1, 2, 0).numpy().astype(np.uint8))
+        raw_images.append(raw_image)
+        image_with_mask = overlay_image_with_mask(images[i].cpu().permute(1, 2, 0).numpy(), pred[i].squeeze(0).cpu().numpy())
+        image_with_mask = Image.fromarray(image_with_mask)
+        segmentation_results.append(image_with_mask)
+    initial_slice_index = 0
+    output_seg = segmentation_results[initial_slice_index]
+    output_raw = raw_images[initial_slice_index]
+    num_slices = len(segmentation_results)
+    initial_size = sizes[initial_slice_index]
+    return output_seg, output_raw, segmentation_results, raw_images, gr.update(maximum=num_slices - 1), sizes, f"Heart volume size: {volume_size} cm^3"
+
+def update_slice(slice_index, segmentation_results_state, raw_images_state, sizes_text):
+    segmentation_results = segmentation_results_state
+    raw_images = raw_images_state
+
+    if segmentation_results is None or raw_images is None:
+        return None, None, ""
+    output_seg = segmentation_results[slice_index]
+    output_raw = raw_images[slice_index]
+
+    return output_seg, output_raw, size_text
+
+def load_example(example):
+    image_file_path = example
+    return inference(image_file_path)
+
+title = "CT Segmentation"
+description = """
+
+<div style="text-align: left; font-weight: bold;">
+    <br>
+    &#x1F32A Note: The current model is run on <span style="color:blue;">CT Segmentation (UW) </span> </p>
+</div>
+"""
+
+article = "The Demo is Run on CT-Seg."
+with gr.Blocks(theme=gr.themes.Soft(), title=title, css=".gradio-container { max-width: 1000px; margin: auto; }") as demo:
+    # add title
+    with gr.Row():
+        gr.Markdown(value="# <span style='color: #6366f1;'>UW CT segmentation</span>", elem_id="title")
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown(value="""
+                Welcome to CT Segmentation, an AI model that segments the thorax and heart out, and computes the volume sizes.
+
+                ## How to Use:
+                0. **Explore Default Examples**: Click on images in the right panel. 
+                1. **Upload Your Image**: something biomedical... but not your lovely pet!
+
+                Click **Segment** and see what CT Seg finds for you!
+                """, 
+                elem_id="instructions")
+            gr.Markdown("## Step 1: Upload CT volume .tif image (Try examples on the right panel)")
+            with gr.Row(equal_height = True):
+                input_image = gr.File(label="Input Image", file_types=[".tif"])
+                # Initially, set the slider maximum to a default value, e.g., 0
+                slice_index_slider = gr.Slider(minimum=0, maximum=0, step=1, label="Slice Index")
+            with gr.Row(equal_height = True):
+                output_raw = gr.Image(label="Processed Image", interactive=False)
+                output_seg = gr.Image(label="Segmentation Results", interactive=False) 
+            with gr.Row():
+                size_text = gr.Textbox(label="Heart volume Size", interactive=False)
+            with gr.Row():
+                button = gr.Button("Segment", interactive=True, variant='primary')
+        with gr.Column(scale=0.5):
+            gr.Markdown("## Click Default Examples")
+            # Initialize state variables
+            segmentation_results_state = gr.State()
+            raw_images_state = gr.State()
+            sizes_state = gr.State()
+            gr.Examples(
+                examples=examples,
+                inputs=[input_image],
+                outputs=[output_seg, output_raw, segmentation_results_state, raw_images_state, slice_index_slider, sizes_state, size_text],
+                fn=load_example,
+                cache_examples=False,
+                examples_per_page=1,
+                run_on_click=True
+            )
+    # Set up the button click
+    button.click(
+        fn=inference,
+        inputs=[input_image],
+        outputs=[output_seg, output_raw, segmentation_results_state, raw_images_state, slice_index_slider, sizes_state, size_text]
+    )
+    # Set up the slider change
+    slice_index_slider.change(
+        fn=update_slice,
+        inputs=[slice_index_slider, segmentation_results_state, raw_images_state, size_text],
+        outputs=[output_seg, output_raw, size_text]
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch(share=True)