Spaces:
Sleeping
Sleeping
File size: 3,527 Bytes
4a65803 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import gradio as gr
import pandas as pd
import numpy as np
import pydicom
import os
from skimage import transform
import torch
from segment_anything import sam_model_registry
import matplotlib.pyplot as plt
from PIL import Image
import io
# Function to load bounding boxes from CSV
def load_bounding_boxes(csv_file):
# Assuming CSV file has columns: 'filename', 'x_min', 'y_min', 'x_max', 'y_max'
df = pd.read_csv(csv_file)
return df
# Function to load DICOM images
def load_dicom_images(folder_path):
images = []
for filename in sorted(os.listdir(folder_path)):
if filename.endswith(".dcm"):
filepath = os.path.join(folder_path, filename)
ds = pydicom.dcmread(filepath)
img = ds.pixel_array
images.append(img)
return np.array(images)
# MedSAM inference function
def medsam_inference(medsam_model, img, box, H, W, target_size):
# Resize image and box to target size
img_resized = transform.resize(img, (target_size, target_size), anti_aliasing=True)
box_resized = np.array(box) * (target_size / np.array([W, H, W, H]))
# Convert image to PyTorch tensor
img_tensor = torch.from_numpy(img_resized).float().unsqueeze(0).unsqueeze(0).to(device) # Add channel and batch dimension
# Model expects box in format (x0, y0, x1, y1)
box_tensor = torch.tensor(box_resized, dtype=torch.float32).unsqueeze(0).to(device) # Add batch dimension
# MedSAM inference
img_embed = medsam_model.image_encoder(img_tensor)
mask = medsam_model.predict(img_embed, box_tensor)
# Post-process mask: resize back to original size
mask_resized = transform.resize(mask[0].cpu().numpy(), (H, W))
return mask_resized
# Function for visualizing images with masks
def visualize(images, masks, box):
fig, ax = plt.subplots(len(images), 2, figsize=(10, 5*len(images)))
for i, (image, mask) in enumerate(zip(images, masks)):
ax[i, 0].imshow(image, cmap='gray')
ax[i, 0].add_patch(plt.Rectangle((box[0], box[1]), box[2]-box[0], box[3]-box[1], edgecolor="red", facecolor="none"))
ax[i, 1].imshow(image, cmap='gray')
ax[i, 1].imshow(mask, alpha=0.5, cmap="jet")
plt.tight_layout()
buf = io.BytesIO()
plt.savefig(buf, format='png')
plt.close(fig)
buf.seek(0)
return buf
# Main function for Gradio app
def process_images(csv_file, dicom_folder, target_size):
bounding_boxes = load_bounding_boxes(csv_file)
dicom_images = load_dicom_images(dicom_folder)
# Initialize MedSAM model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
medsam_model = sam_model_registry['your_model_version'](checkpoint='path_to_your_checkpoint')
medsam_model = medsam_model.to(device)
medsam_model.eval()
masks = []
for index, row in bounding_boxes.iterrows():
if index >= len(dicom_images):
continue # Skip if the index exceeds the number of images
image = dicom_images[index]
H, W = image.shape
box = [row['x_min'], row['y_min'], row['x_max'], row['y_max']]
mask = medsam_inference(medsam_model, image, box, H, W, target_size)
masks.append(mask)
visualizations = visualize(dicom_images, masks, box)
return visualizations, np.array(masks)
# Set up Gradio interface
iface = gr.Interface(
fn=process_images,
inputs=[gr.inputs.File(type="file"), gr.inputs.Directory()],
outputs=[gr.outputs.Image(type="plot"), gr.outputs.File(type="numpy")]
)
iface.launch()
|