Merge pull request #12 from bioimage-io/ray-sam

Run SAM computations with RAY

.gitignore

@@ -14,6 +14,7 @@ data/
 .DS_Store
 *.onnx
 *.pt
+*.pth
 *tif
 *zip
 visualize_annotation.py

Dockerfile

@@ -31,19 +31,16 @@ COPY ./requirements-sam.txt /app/requirements-sam.txt
 RUN pip install -r /app/requirements-sam.txt
 
 # Copy the python script to the docker environment
-COPY ./bioimageio_colab/register_sam_service.py /app/register_sam_service.py
+COPY ./bioimageio_colab /app/bioimageio_colab
 
-# Copy the start service script
-COPY ./scripts/start_service.sh /app/start_service.sh
+# Create cache directory for models
+RUN mkdir -p /app/.model_cache
 
 # Change ownership of the application directory to the non-root user
 RUN chown -R bioimageio_colab:bioimageio_colab /app/
 
-# Make the start script executable
-RUN chmod +x /app/start_service.sh
-
 # Switch to the non-root user
 USER bioimageio_colab
 
 # Use the start script as the entrypoint and forward arguments
-ENTRYPOINT ["python", "register_sam_service.py"]
+ENTRYPOINT ["python", "-m", "bioimageio_colab"]

bioimageio_colab/__main__.py

@@ -1,5 +1,42 @@
-def main():
-  print("Welcome to the bioimageio-colab")
+import argparse
+import asyncio
 
-if __name__ == '__main__':
-    main()
+from bioimageio_colab.register_sam_service import register_service
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Register SAM annotation service on BioImageIO Colab workspace."
+    )
+    parser.add_argument(
+        "--server_url",
+        default="https://hypha.aicell.io",
+        help="URL of the Hypha server",
+    )
+    parser.add_argument(
+        "--workspace_name", default="bioimageio-colab", help="Name of the workspace"
+    )
+    parser.add_argument(
+        "--service_id",
+        default="microsam",
+        help="Service ID for registering the service",
+    )
+    parser.add_argument(
+        "--token",
+        default=None,
+        help="Workspace token for connecting to the Hypha server",
+    )
+    parser.add_argument(
+        "--cache_dir",
+        default="./.model_cache",
+        help="Directory for caching the models",
+    )
+    parser.add_argument(
+        "--ray_address",
+        default=None,
+        help="Address of the Ray cluster for running SAM",
+    )
+    args = parser.parse_args()
+
+    loop = asyncio.get_event_loop()
+    loop.create_task(register_service(args=args))
+    loop.run_forever()

bioimageio_colab/create_workspace.py

@@ -1,7 +1,12 @@
+"""
+Hypha now supports creating workspaces from the Hypha Dashboard. This script is no longer needed.
+"""
+
 import argparse
-from hypha_rpc import connect_to_server, login
 import asyncio
 
+from hypha_rpc import connect_to_server, login
+
 
 async def create_workspace_token(args):
     # Get a user login token
@@ -64,6 +69,7 @@ async def create_workspace_token(args):
         for service in services:
             print(f"- {service['name']}")
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Create the BioImageIO Colab workspace."