def check_device(gpu_number=0, callback_log=None):
""" Check on which device (GPU or CPU) CellPose will be running.
Parameters
----------
gpu_number : int, optional
Index of GPU to be tested, by default 0
"""
use_gpu = utils.use_gpu(gpu_number=gpu_number)
if use_gpu:
log_message('CellPose will be USING GPU', callback_fun=callback_log)
device = mx.gpu()
else:
log_message('CellPose will be USING CPU', callback_fun=callback_log)
device = mx.cpu()
return device
def initialize_model(init_model, model_type, which_model):
# initialize models
# returns a dictonary of models that can be called
mds = {}
# UNet Cell Profiler
if model_type == "Cellprofiler_UNet":
option_dict_conv, option_dict_bn = init_model["UNetSettings"]
os.environ["KERAS_BACKEND"] = "tensorflow"
if which_model == "UNet_CP001":
model = unet_initialize(init_model["UNetShape"],
option_dict_conv,
option_dict_bn,
init_model["UNet_model_file_CP01"],
automated_shape_adjustment=True)
mds["UNet_CP001"] = model
# Stardist
elif model_type == "StarDist":
if which_model == "SD_2D_dsb2018":
model = StarDist2D(None,
name='2D_dsb2018',
basedir=init_model["basedir_StarDist"])
mds["SD_2D_dsb2018"] = model
# Cellpose
elif model_type == "Cellpose":
# check if GPU working, and if so use it
use_gpu = cp_utils.use_gpu()
if use_gpu:
device = mx.gpu()
print("GPU found")
else:
device = mx.cpu()
print("CPU only")
if which_model == "CP_nuclei":
model = cp_models.Cellpose(device, model_type="nuclei")
mds["CP_nuclei"] = model
if which_model == "CP_cyto":
model = cp_models.Cellpose(device, model_type="cyto")
mds["CP_cyto"] = model
print("Model_keys", mds.keys())
return mds
image_names = get_image_files(args.dir)
imn = []
for im in image_names:
imfile = os.path.splitext(im)[0]
if len(imfile) > len(args.mask_filter):
if imfile[-len(args.mask_filter):] != args.mask_filter:
imn.append(im)
else:
imn.append(im)
image_names = imn
nimg = len(image_names)
#nimg =
images = [skimage.io.imread(image_names[n]) for n in range(nimg)]
if args.use_gpu:
use_gpu = utils.use_gpu()
if use_gpu:
device = mx.gpu()
else:
device = mx.cpu()
print('>>>> using %s' % (['CPU', 'GPU'][use_gpu]))
if args.pretrained_model == 'cyto' or args.pretrained_model == 'nuclei':
cp_path = os.path.dirname(os.path.realpath(__file__))
if not args.train:
model_list = [
'%s_%d' % (args.pretrained_model, i) for i in range(4)
]
cpmodel_path = [
os.path.abspath(
os.path.join(cp_path, 'models/', model_list[i]))