Spaces:
Running
Running
| # Code copied and modified from https://huggingface.co/spaces/BAAI/SegVol/blob/main/utils.py | |
| from pathlib import Path | |
| import matplotlib as mpl | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import SimpleITK as sitk | |
| import torch | |
| from mrsegmentator import inference | |
| from mrsegmentator.utils import add_postfix | |
| from PIL import Image | |
| from scipy import ndimage | |
| import streamlit as st | |
| initial_rectangle = { | |
| "version": "4.4.0", | |
| "objects": [ | |
| { | |
| "type": "rect", | |
| "version": "4.4.0", | |
| "originX": "left", | |
| "originY": "top", | |
| "left": 50, | |
| "top": 50, | |
| "width": 100, | |
| "height": 100, | |
| "fill": "rgba(255, 165, 0, 0.3)", | |
| "stroke": "#2909F1", | |
| "strokeWidth": 3, | |
| "strokeDashArray": None, | |
| "strokeLineCap": "butt", | |
| "strokeDashOffset": 0, | |
| "strokeLineJoin": "miter", | |
| "strokeUniform": True, | |
| "strokeMiterLimit": 4, | |
| "scaleX": 1, | |
| "scaleY": 1, | |
| "angle": 0, | |
| "flipX": False, | |
| "flipY": False, | |
| "opacity": 1, | |
| "shadow": None, | |
| "visible": True, | |
| "backgroundColor": "", | |
| "fillRule": "nonzero", | |
| "paintFirst": "fill", | |
| "globalCompositeOperation": "source-over", | |
| "skewX": 0, | |
| "skewY": 0, | |
| "rx": 0, | |
| "ry": 0, | |
| } | |
| ], | |
| } | |
| def run(tmpdirname): | |
| if st.session_state.option is not None: | |
| image = Path(__file__).parent / str(st.session_state.option) | |
| inference.infer([image], tmpdirname, st.session_state.folds, split_level=1) | |
| seg_name = add_postfix(image.name, "seg") | |
| preds_path = tmpdirname + "/" + seg_name | |
| st.session_state.preds_3D = read_image(preds_path) | |
| st.session_state.preds_3D_ori = sitk.ReadImage(preds_path) | |
| def reflect_box_into_model(box_3d): | |
| z1, y1, x1, z2, y2, x2 = box_3d | |
| x1_prompt = int(x1 * 256.0 / 325.0) | |
| y1_prompt = int(y1 * 256.0 / 325.0) | |
| z1_prompt = int(z1 * 32.0 / 325.0) | |
| x2_prompt = int(x2 * 256.0 / 325.0) | |
| y2_prompt = int(y2 * 256.0 / 325.0) | |
| z2_prompt = int(z2 * 32.0 / 325.0) | |
| return torch.tensor(np.array([z1_prompt, y1_prompt, x1_prompt, z2_prompt, y2_prompt, x2_prompt])) | |
| def reflect_json_data_to_3D_box(json_data, view): | |
| if view == "xy": | |
| st.session_state.rectangle_3Dbox[1] = json_data["objects"][0]["top"] | |
| st.session_state.rectangle_3Dbox[2] = json_data["objects"][0]["left"] | |
| st.session_state.rectangle_3Dbox[4] = ( | |
| json_data["objects"][0]["top"] + json_data["objects"][0]["height"] * json_data["objects"][0]["scaleY"] | |
| ) | |
| st.session_state.rectangle_3Dbox[5] = ( | |
| json_data["objects"][0]["left"] + json_data["objects"][0]["width"] * json_data["objects"][0]["scaleX"] | |
| ) | |
| print(st.session_state.rectangle_3Dbox) | |
| def make_fig(image, preds, px_range=(10, 400), transparency=0.5): | |
| fig, ax = plt.subplots(1, 1, figsize=(4, 4)) | |
| image_slice = image.clip(*px_range) | |
| ax.imshow( | |
| image_slice, | |
| cmap="Greys_r", | |
| vmin=px_range[0], | |
| vmax=px_range[1], | |
| ) | |
| if preds is not None: | |
| image_slice = np.array(preds) | |
| alpha = np.zeros(image_slice.shape) | |
| alpha[image_slice > 0.1] = transparency | |
| ax.imshow( | |
| image_slice, | |
| cmap="jet", | |
| alpha=alpha, | |
| vmin=0, | |
| vmax=40, | |
| ) | |
| # plot edges | |
| edge_slice = np.zeros(image_slice.shape, dtype=int) | |
| for i in np.unique(image_slice): | |
| _slice = image_slice.copy() | |
| _slice[_slice != i] = 0 | |
| edges = ndimage.laplace(_slice) | |
| edge_slice[edges != 0] = i | |
| cmap = mpl.cm.jet(np.linspace(0, 1, int(preds.max()))) | |
| cmap -= 0.4 | |
| cmap = cmap.clip(0, 1) | |
| cmap = mpl.colors.ListedColormap(cmap) | |
| alpha = np.zeros(edge_slice.shape) | |
| alpha[edge_slice > 0.01] = 0.9 | |
| ax.imshow( | |
| edge_slice, | |
| alpha=alpha, | |
| cmap=cmap, | |
| vmin=0, | |
| vmax=40, | |
| ) | |
| plt.axis("off") | |
| ax.set_xticks([]) | |
| ax.set_yticks([]) | |
| fig.canvas.draw() | |
| # transform to image | |
| return Image.frombytes("RGB", fig.canvas.get_width_height(), fig.canvas.tostring_rgb()) | |
| ####################################### | |
| def make_isotropic(image, interpolator=sitk.sitkLinear, spacing=None): | |
| """ | |
| Many file formats (e.g. jpg, png,...) expect the pixels to be isotropic, same | |
| spacing for all axes. Saving non-isotropic data in these formats will result in | |
| distorted images. This function makes an image isotropic via resampling, if needed. | |
| Args: | |
| image (SimpleITK.Image): Input image. | |
| interpolator: By default the function uses a linear interpolator. For | |
| label images one should use the sitkNearestNeighbor interpolator | |
| so as not to introduce non-existant labels. | |
| spacing (float): Desired spacing. If none given then use the smallest spacing from | |
| the original image. | |
| Returns: | |
| SimpleITK.Image with isotropic spacing which occupies the same region in space as | |
| the input image. | |
| """ | |
| original_spacing = image.GetSpacing() | |
| # Image is already isotropic, just return a copy. | |
| if all(spc == original_spacing[0] for spc in original_spacing): | |
| return sitk.Image(image) | |
| # Make image isotropic via resampling. | |
| original_size = image.GetSize() | |
| if spacing is None: | |
| spacing = min(original_spacing) | |
| new_spacing = [spacing] * image.GetDimension() | |
| new_size = [int(round(osz * ospc / spacing)) for osz, ospc in zip(original_size, original_spacing)] | |
| return sitk.Resample( | |
| image, | |
| new_size, | |
| sitk.Transform(), | |
| interpolator, | |
| image.GetOrigin(), | |
| new_spacing, | |
| image.GetDirection(), | |
| 0, # default pixel value | |
| image.GetPixelID(), | |
| ) | |
| def read_image(path): | |
| img = sitk.ReadImage(path) | |
| img = sitk.DICOMOrient(img, "LPS") | |
| img = make_isotropic(img) | |
| img = sitk.GetArrayFromImage(img) | |
| return img | |