Add foreign env runner

Author: gnodar01

CP5/active_plugins/cpforeign/server.py

@@ -42,19 +42,19 @@ def __init__(self, port, domain='*', protocol='tcp', wait_for_handshake=True):
         self._context = zmq.Context()
         self._server_socket = self._context.socket(zmq.PAIR)
         self._server_socket.bind(f"{protocol}://{domain}:{port}")
-        self._logger.info("launched on", self._server_socket.getsockopt(zmq.LAST_ENDPOINT))
+        self._logger.info(f"launched on {self._server_socket.getsockopt(zmq.LAST_ENDPOINT)}")
 
         if wait_for_handshake:
             self.wait_for_handshake()
 
     def wait_for_handshake(self):
-        self._logger.debug("waiting for handshake from cllient")
+        self._logger.debug("waiting for handshake from client")
         client_hello = self._server_socket.recv_string()
         if client_hello == HELLO:
             self._logger.debug("received correct handshake")
             _send_ack(self._server_socket, self._logger, subject="handshake")
         else:
-            self._logger.debug("received incorrect handshake", client_hello)
+            self._logger.debug(f"received incorrect handshake {client_hello}")
             self._logger.debug("sending deny")
             self._server_socket.send_string(DENY)
             raise ForeignToolError("server received incorrect handshake")
@@ -65,12 +65,12 @@ def _serve_image(self, image_data):
         """
         header = np.lib.format.header_data_from_array_1_0(image_data)
 
-        self._logger.debug("sending header", header, "waiting for acknowledgement")
+        self._logger.debug(f"sending header {header} waiting for acknowledgement")
         self._server_socket.send_json(header)
 
         ack = _receive_ack(self._server_socket, self._logger, subject="header")
 
-        self._logger.debug("sending image data", image_data.shape, "waiting for acknowledgement")
+        self._logger.debug(f"sending image data {image_data.shape} waiting for acknowledgement")
 
         self._server_socket.send(image_data, copy=False)
 
@@ -87,7 +87,7 @@ def _serve_image(self, image_data):
         self._logger.debug("parsing label data")
         labels = np.frombuffer(label_bytes, dtype=labels_header['descr'])
         labels.shape = labels_header['shape']
-        self._logger.debug("parse label data", labels.shape)
+        self._logger.debug(f"parse label data of shape {labels.shape}")
 
         _send_ack(self._server_socket, self._logger, subject="return data")
 
@@ -100,7 +100,7 @@ def serve_one_image(self, image_data):
         return self._serve_image(image_data)
 
 class ForeignToolClient(object):
-    def __init__(self, port, domain='*', protocol='tcp', do_handshake=True, cb=None):
+    def __init__(self, port, domain='localhost', protocol='tcp', do_handshake=True, cb=None):
         """
         Connect to a server on the given port.
         """
@@ -109,20 +109,20 @@ def __init__(self, port, domain='*', protocol='tcp', do_handshake=True, cb=None)
         self._context = zmq.Context()
         self._client_socket = self._context.socket(zmq.PAIR)
         self._client_socket.connect(f"{protocol}://{domain}:{port}")
-        self._logger.info("connected to", self._client_socket.getsockopt(zmq.LAST_ENDPOINT))
+        self._logger.info(f"connected to {self._client_socket.getsockopt(zmq.LAST_ENDPOINT)}")
 
         if cb:
             self.register_cb(cb)
 
         if do_handshake:
             self.do_handshake()
 
-        def do_handshake(self):
-            """
-            Handshake with the server.
-            """
-            self.client_socket.send_string(HELLO)
-            response = _receive_ack(self.client_socket, self.logger, subject="handshake")
+    def do_handshake(self):
+        """
+        Handshake with the server.
+        """
+        self._client_socket.send_string(HELLO)
+        response = _receive_ack(self._client_socket, self._logger, subject="handshake")
 
     def register_cb(self, cb):
         """
@@ -142,7 +142,7 @@ def _receive_image(self):
         Receive an image from the server.
         """
         header = self._client_socket.recv_json()
-        self._logger.debug("received header", header)
+        self._logger.debug(f"received header {header}")
 
         _send_ack(self._client_socket, self._logger, subject="header")
 
@@ -152,9 +152,8 @@ def _receive_image(self):
         self._logger.debug("parsing image data")
         buf = memoryview(im_bytes)
         im = np.frombuffer(buf, dtype=header['descr'])
-        im = (im * 255).astype(np.uint8)
         im.shape = header['shape']
-        self._logger.debug("parsed image data", im.shape)
+        self._logger.debug(f"parsed image data {im.shape}")
 
         _send_ack(self._client_socket, self._logger, subject="image data")
 
@@ -163,7 +162,7 @@ def _receive_image(self):
         self._logger.debug("executed callback")
 
         return_header = np.lib.format.header_data_from_array_1_0(return_data)
-        self._logger.debug("returning header", return_header)
+        self._logger.debug(f"returning header {return_header}")
         self._client_socket.send_json(return_header)
 
         ack = _receive_ack(self._client_socket, self._logger, subject="return header")

CP5/active_plugins/cpforeign/thresh.py

@@ -0,0 +1,54 @@
+import logging
+import numpy as np
+import skimage as ski
+import scipy as sp
+
+from server import ForeignToolClient
+
+logger = logging.getLogger(__name__)
+
+
+def run(image_data, image_header):
+    im = (image_data * 255).astype(np.uint8)
+
+    markers = np.zeros_like(im, dtype=np.uint8)
+    IDK = 0
+    BG = 1
+    FG = 2
+    markers[im < 30] = BG
+    markers[im > 50] = FG
+    # rest = IDK
+
+    elevation_map = ski.filters.sobel(im)
+    segmentation = ski.segmentation.watershed(elevation_map, markers)
+    segmentation = sp.ndimage.binary_fill_holes(segmentation - 1)
+
+    labels, _ = sp.ndimage.label(segmentation)
+
+    # remove small objects
+    sizes = np.bincount(labels.ravel())
+    mask_sizes = sizes > 20
+    mask_sizes[0] = 0
+    segmentation = mask_sizes[labels]
+
+    labels, _ = sp.ndimage.label(segmentation)
+    
+    ski.io.imsave("/Users/ngogober/Desktop/thresh.tif", labels)
+
+    return labels
+
+
+def main():
+    client = ForeignToolClient(7878, cb=run)
+    client.receive_images()
+
+if __name__ == "__main__":
+    # init logging
+    logging.root.setLevel(logging.DEBUG)
+    stream_handler = logging.StreamHandler()
+    fmt = logging.Formatter("  [%(process)d|%(levelno)s] %(name)s::%(funcName)s: %(message)s")
+    stream_handler.setFormatter(fmt)
+    logging.root.addHandler(stream_handler)
+
+    logger.debug("Starting thresh.py")
+    main()

CP5/active_plugins/cpforeign/zmq_server.ipynb

@@ -0,0 +1,172 @@
+import shlex
+import sys
+import os
+import re
+import subprocess
+import threading
+import logging
+
+from cellprofiler_core.module.image_segmentation import ImageSegmentation
+from cellprofiler_core.setting.text import Text
+from cellprofiler_core.setting.text import Filename
+from cellprofiler_core.setting.text import Integer
+from cellprofiler_core.object import Objects
+
+from cpforeign.server import ForeignToolServer
+
+LOGGER = logging.getLogger(__name__)
+
+HELLO = "Hello"
+ACK = "Acknowledge"
+DENIED = "Denied"
+
+__doc__ = """\
+RunForeignEnv
+============
+
+**RunForeign** runs a foreign tool, in a foreign (conda) environment, via sockets.
+
+
+Assumes there is a client up and running.
+
+|
+
+============ ============ ===============
+Supports 2D? Supports 3D? Respects masks?
+============ ============ ===============
+YES          NO            YES
+============ ============ ===============
+
+"""
+
+def _run_logger(workR):
+    # this thread shuts itself down by reading from worker's stdout
+    # which either reads content from stdout or blocks until it can do so
+    # when the worker is shut down, empty byte string is returned continuously
+    # which evaluates as None so the break is hit
+    # I don't really like this approach; we should just shut it down with the other
+    # threads explicitly
+    while True:
+        try:
+            print('reading')
+            line = workR.stdout.readline()
+            if (type(line) == bytes):
+                line = line.decode("utf-8")
+            if not line:
+                break
+            log_msg_match = re.match(fr"{workR.pid}\|(10|20|30|40|50)\|(.*)", line)
+            if log_msg_match:
+                levelno = int(log_msg_match.group(1))
+                msg = log_msg_match.group(2)
+            else:
+                levelno = 20
+                msg = line
+
+            LOGGER.log(levelno, "\n\r  [Worker (%d)] %s", workR.pid, msg.rstrip())
+
+        except Exception as e:
+            LOGGER.exception(e)
+            break
+
+
+class RunForeignEnv(ImageSegmentation):
+    category = "Object Processing"
+
+    module_name = "RunForeignEnv"
+
+    variable_revision_number = 1
+
+    def create_settings(self):
+        super().create_settings()
+        
+        self._server = None
+        self._client_launched = False
+
+        self.server_port = Integer(
+            text="Server port number",
+            value=7878,
+            minval=0,
+            doc="""\
+The port number which the server is listening on. The server must be launched manually first.
+""",
+        )
+        
+        self.env_name = Text(text="Conda environment name", value="foreign-thresh")
+        
+        self.algo_path = Filename(text="Algorithm path", value="/Users/ngogober/Developer/CellProfiler/CellProfiler-plugins/CP5/active_plugins/cpforeign/thresh.py")
+
+    def settings(self):
+        return super().settings() + [self.server_port, self.env_name, self.algo_path]
+    
+    # ImageSegmentation defines this so we have to overide it
+    def visible_settings(self):
+        return self.settings()
+    
+    # ImageSegmentation defines this so we have to overide it
+    def volumetric(self):
+        return False
+
+    def prepare_run(self, workspace):
+
+        LOGGER.debug(">>> Preparing run")
+        if not self._server:
+            LOGGER.debug(">>> Initializing server")
+            self._server = ForeignToolServer(self.server_port.value, wait_for_handshake=False)
+
+        if not self._client_launched:
+            LOGGER.debug(">>> Launching client")
+            command = f"conda run --no-capture-output -n {self.env_name.value} python {self.algo_path.value}"
+            args = shlex.split(command)
+            env = os.environ.copy()
+            env["PYTHONUNBUFFERED"] = "1"
+            self._client_proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=sys.stdout, bufsize=1, universal_newlines=True, env=env, close_fds=False)
+            #self._client_thread = threading.Thread(target=_run_logger, args=(self._client_proc,), name="foreign client stdout logger thread")
+            #self._client_thread.start()
+
+            self._client_launched = True
+            self._server.wait_for_handshake()
+
+        return True
+
+    def post_run(self, workspace):
+        if self._client_launched:
+            LOGGER.debug(">>> Shuttding down client")
+            #self._client_thread.join()
+            self._client_proc.terminate()
+
+    def run(self, workspace):
+        # TODO: is this supposed to not run in test mode? because it doesn't...
+        self.prepare_run(workspace)
+
+        x_name = self.x_name.value
+
+        y_name = self.y_name.value
+
+        images = workspace.image_set
+
+        x = images.get_image(x_name)
+
+        dimensions = x.dimensions
+
+        x_data = x.pixel_data
+
+        y_data = self._server.serve_one_image(x_data)
+
+        y = Objects()
+
+        y.segmented = y_data
+
+        y.parent_image = x.parent_image
+
+        objects = workspace.object_set
+
+        objects.add_objects(y, y_name)
+
+        self.add_measurements(workspace)
+
+        if self.show_window:
+            workspace.display_data.x_data = x_data
+
+            workspace.display_data.y_data = y_data
+
+            workspace.display_data.dimensions = dimensions