infer.py

from segment_anything import SamPredictor, sam_model_registry
from get_bbox import get_bounding_box
import SimpleITK as sitk
import numpy as np
import os
from MedSAM_Inference import infer

img_path = '/path/to/imagesTr'
label_path = '/path/to/labelsTr'
img_names = sorted(os.listdir(img_path))
label_names = sorted(os.listdir(label_path))

for img_name, label_name in zip(img_names, label_names):
    img = sitk.ReadImage(os.path.join(img_path, img_name))
    label = sitk.ReadImage(os.path.join(label_path, label_name))
    img_arr = sitk.GetArrayFromImage(img)
    label_arr = sitk.GetArrayFromImage(label)
    img_arr = np.transpose(img_arr, (2, 1, 0))
    label_arr = np.transpose(label_arr, (2, 1, 0))
    x_len,y_len,z_len = img_arr.shape
    bbox = get_bounding_box(label_arr)
    masks = []
    for z in range(0, bbox["z_min"]):
        masks.append(np.zeros((1, x_len, y_len)))
    for z in range(bbox["z_min"], bbox["z_max"] + 1):
        input_slice = img_arr[..., z]
        input_slice = np.expand_dims(input_slice, axis=-1)
        input_slice = (input_slice - input_slice.min())/(input_slice.max() - input_slice.min()+1)*255
        input_slice = input_slice.astype(np.uint8)
        input_slice = np.repeat(input_slice, 3, axis=-1)
        input_box = str([bbox["y_min"], bbox["x_min"], bbox["y_max"], bbox["x_max"]])
        mask = infer(input_slice, input_box)
        mask = np.expand_dims(mask, axis=0)
        masks.append(mask)
    for z in range(bbox["z_max"] + 1, z_len):
        masks.append(np.zeros((1, x_len, y_len)))
    masks = np.stack(masks, axis=0)
    masks = masks.squeeze(axis=1)
    masks = np.transpose(masks, (0, 2, 1))
    masks = masks.astype(np.uint8)
    masks = sitk.GetImageFromArray(masks)
    masks.CopyInformation(img)
    output_path = '/path/to/medsam_infer/'
    output_name = label_name
    sitk.WriteImage(masks, os.path.join(output_path, output_name))