# Package required to interact with Cellpose
%pip install cellpose==2.2.1

# Package required to interact with IDR or OMERO
%pip install omero-py==5.13.1

# Package to convert masks stored in OMERO/IDR into numpy arrays
%pip install omero-cli-zarr==0.5.2

image_id = 6001247

image = conn.getObject("Image", image_id)
print(image.getName())

print("Size X: %s, Size Y: %s" % (image.getSizeX(), image.getSizeY()))

data = load_numpy_array(image)

print(data.shape)

channels = [[0, 1]]
t = 0
z = int(image.getPrimaryPixels().getSizeZ() / 2)
cellpose_masks, flows, styles, diams = model.eval(data[t, :, z, :, :], diameter=None, channels=channels)

roi_service = conn.getRoiService()
result = roi_service.findByImage(image_id, None)

shapes = []
for roi in result.rois:
    shapes.append(roi.copyShapes())

from omero_zarr import masks
dims = (image.getSizeT(), image.getSizeC(), image.getSizeZ(), image.getSizeY(), image.getSizeX())
saver = masks.MaskSaver(None, image, numpy.int64)
labels, fillColors, properties = saver.masks_to_labels(shapes, mask_shape=dims)

import matplotlib.pyplot as plt
%matplotlib inline
c = 1
t = 0
fig = plt.figure(figsize=(10, 10))
# Show the original labels from IDR
sub1 = plt.subplot(121)
sub1.title.set_text('Original labels IDR')
plt.imshow(labels[t, c, z, :, :], cmap='gray', alpha=0.5)
# Show the Cellpose labels
sub2 = plt.subplot(122)
sub2.title.set_text('Cellpose labels this notebook')
plt.imshow(cellpose_masks, cmap='jet', alpha=0.5)
fig2 = plt.figure(figsize=(8.5, 8.5))
# Show the original image
sub3 = plt.subplot(121)
sub3.title.set_text("Original unsegmented image")
plt.imshow(data[t, c, z, :, :], cmap='jet')
# Show the Cellpose labels
sub4 = plt.subplot(122)
sub4.title.set_text('Cellpose labels this notebook')
plt.imshow(cellpose_masks, cmap='jet', alpha=0.5)
plt.tight_layout()
fig.canvas.flush_events()

from cellpose import io, utils
outlines = utils.outlines_list(cellpose_masks)
io.outlines_to_text(str(image_id), outlines)

def create_roi(img, shapes):
    # create an ROI, link it to Image
    roi = omero.model.RoiI()
    # use the omero.model.ImageI that underlies the 'image' wrapper
    roi.setImage(img._obj)
    for shape in shapes:
        roi.addShape(shape)
    return roi

# Convert roi into ome rois
import re
import omero
from omero.rtypes import rdouble, rint, rstring
polygons = []
with open(str(image_id) + "_cp_outlines.txt", "r") as text_file:
    for line in text_file:
        points = re.sub(r',([^,]*),', r',\1, ', line)
        polygon = omero.model.PolygonI()
        polygon.theZ = rint(z)
        polygon.theT = rint(t)
        polygon.points = rstring(points)
        polygons.append(polygon)
roi = create_roi(image, polygons)

# Save the ROI when using an OMERO server
if "idr" not in HOST:
    conn.getUpdateService().saveAndReturnObject(roi)

def disconnect(conn):
    """
    Disconnect from an OMERO server
    :param conn: The BlitzGateway
    """
    conn.close()

disconnect(conn)