adds png plotting to each step (#148)

* adds png plotting to each step

* adds step to plot original reconstruction

pipeline/commands/command_utils.py

@@ -0,0 +1,67 @@
+import os
+
+import matplotlib.pyplot as plt
+import boto3
+
+from neuron_morphology.morphology import Morphology
+from neuron_morphology.swc_io import (morphology_from_swc, morphology_to_swc)
+
+
+s3 = boto3.client("s3")
+
+
+def collect_morphology(bucket: str,
+                       swc_key: str):
+    swc_response = s3.get_object(Bucket=bucket,
+                                 Key=swc_key,
+                                 )
+    return morphology_from_swc(swc_response["Body"])
+
+
+def morphology_to_s3(bucket: str, key: str, morphology: Morphology):
+    """Store morphology in an s3 bucket as swc
+
+    Parameters
+    ----------
+    bucket : name of the bucket in which to store this dataset
+    key : at which to store the dataset
+    morphology : The morphology to store
+
+    Returns
+    -------
+    key : the argued key
+
+    Notes
+    -----
+    we must write morphology to "disk" temporarily and send that.
+    """
+    tmp_path = key.split("/")[-1]
+    morphology_to_swc(morphology, tmp_path)
+
+    s3.upload_file(Filename=tmp_path, Bucket=bucket, Key=key)
+    os.remove(tmp_path)
+    return key
+
+
+def morphology_png_to_s3(bucket: str,
+                         key: str,
+                         morphology: Morphology):
+
+    tmp_path = key.split("/")[-1]
+
+    nodes = morphology.nodes()
+    x = [node['x'] for node in nodes]
+    y = [node['y'] for node in nodes]
+    z = [node['z'] for node in nodes]
+
+    fig, ax = plt.subplots(1, 2)
+    ax[0].scatter(x, y, s=0.1)
+    ax[0].set_title('x-y view')
+    ax[1].scatter(z, y, s=0.1)
+    ax[1].set_title('z-y view')
+    fig.suptitle(tmp_path[:-4], fontsize=16)
+    fig.savefig(tmp_path)
+
+    s3.upload_file(Filename=tmp_path, Bucket=bucket, Key=key)
+    os.remove(tmp_path)
+    return key

pipeline/commands/plot_original.py

@@ -0,0 +1,40 @@
+import os
+import json
+from typing import Optional
+
+import boto3
+
+from command_utils import collect_morphology, morphology_png_to_s3
+from harness import step_fns_ecs_harness
+
+s3 = boto3.client("s3")
+
+
+def plot_original(token: Optional[str] = None):
+    reconstruction_id = os.environ["RECONSTRUCTION_ID"]
+    working_bucket = os.environ["WORKING_BUCKET"]
+    run_prefix = os.environ["RUN_PREFIX"]
+
+    input_json_key = f"{run_prefix}/{reconstruction_id}.json"
+    input_json_response = s3.get_object(Bucket=working_bucket,
+                                        Key=input_json_key
+                                        )
+
+    input_data = json.load(input_json_response["Body"])
+
+    swc_key = f"{run_prefix}/{input_data['swc_file']}"
+    morphology = collect_morphology(working_bucket, swc_key)
+    return {
+            "original_png_key": morphology_png_to_s3(
+                working_bucket,
+                f"{run_prefix}/original_morphology.png",
+                morphology)
+            }
+
+
+def main():
+    step_fns_ecs_harness(plot_original)
+
+
+if __name__ == "__main__":
+    main()

pipeline/commands/scale_correction.py

@@ -5,17 +5,17 @@
 from neuron_morphology.swc_io import morphology_from_swc
 from typing import Dict, Any
 from neuron_morphology.morphology import Morphology
-from neuron_morphology.swc_io import morphology_to_swc
 
+from command_utils import morphology_to_s3, morphology_png_to_s3
 from harness import step_fns_ecs_harness
 
 s3 = boto3.client("s3")
 
 
 def collect_inputs(
-    working_bucket: str,
-    run_prefix: str,
-    reconstruction_id: int) -> Dict[str,Any]:
+        working_bucket: str,
+        run_prefix: str,
+        reconstruction_id: int) -> Dict[str, Any]:
     """
     Gather from AWS the inputs  required to run the scale correction module.
 
@@ -48,14 +48,13 @@ def collect_inputs(
     swc_key = f"{run_prefix}/{input_data['swc_file']}"
     soma_marker_key = f"{run_prefix}/{input_data['marker_file']}"
 
-
     swc_response = s3.get_object(Bucket=working_bucket,
-                                      Key=swc_key,
-                                      )
+                                 Key=swc_key,
+                                 )
 
     soma_marker_response = s3.get_object(Bucket=working_bucket,
-                                      Key=soma_marker_key,
-                                      )
+                                         Key=soma_marker_key,
+                                         )
 
     morphology = morphology_from_swc(swc_response["Body"])
 
@@ -69,31 +68,6 @@ def collect_inputs(
     }
 
 
-def morphology_to_s3(bucket: str, key: str, morphology: Morphology):
-    """Store morphology in an s3 bucket as swc
-
-    Parameters
-    ----------
-    bucket : name of the bucket in which to store this dataset
-    key : at which to store the dataset
-    morphology : The morphology to store
-
-    Returns
-    -------
-    key : the argued key
-
-    Notes
-    -----
-    we must write morphology to "disk" temporarily and send that.
-    """
-    tmp_path = key.split("/")[-1]
-    morphology_to_swc(morphology, tmp_path)
-
-    s3.upload_file(Filename=tmp_path, Bucket=bucket, Key=key)
-    os.remove(tmp_path)
-    return key
-
-
 def put_outputs(
         bucket: str,
         prefix: str,
@@ -116,7 +90,8 @@ def put_outputs(
     -------
     Outputs to send back to step functions. These are:
         scale_transform : Dict[str,float]
-        scaled_morph_key : s3 key of the scaled_morphology
+        scaled_swc_key : s3 key of the scaled_morphology
+        scaled_png_key: s3 key of the scaled morphology png
     """
     return {
         "scale_transform": transform,
@@ -126,6 +101,11 @@ def put_outputs(
             f"{prefix}/scaled_morphology.swc",
             morphology
         ),
+        "scaled_png_key": morphology_png_to_s3(
+            bucket,
+            f"{prefix}/scaled_morphology.png",
+            morphology
+        )
     }
 
 

pipeline/commands/tilt.py

@@ -11,7 +11,7 @@
     import (AffineTransform)
 
 from harness import step_fns_ecs_harness
-from scale_correction import morphology_to_s3
+from command_utils import morphology_to_s3, morphology_png_to_s3
 
 
 s3 = boto3.client("s3")
@@ -67,7 +67,12 @@ def put_outputs(bucket,
         'tilt_correction': tilt_correction,
         'tilt_transform': tilt_transform.to_dict(),
         'tilt_swc_key': morphology_to_s3(
-            bucket, f"{prefix}/tilt_morphology.swc", tilted_morphology)
+            bucket, f"{prefix}/tilt_morphology.swc", tilted_morphology),
+        'tilt_png_key': morphology_png_to_s3(
+            bucket,
+            f"{prefix}/tilt_morphology.png",
+            tilted_morphology
+        )
     }
 
 

pipeline/commands/upright_transform.py

@@ -14,7 +14,7 @@
 from neuron_morphology.morphology import Morphology
 
 from harness import step_fns_ecs_harness
-from scale_correction import morphology_to_s3
+from command_utils import morphology_to_s3, morphology_png_to_s3
 
 
 s3 = boto3.client("s3")
@@ -75,14 +75,18 @@ def put_outputs(
     Outputs to send back to step functions. These are:
         upright_transform_dict : the upright transform matrix
         upright_angle : the upright angle
-        upright_morphology_key : s3 key of the upright transformed morphology 
+        upright_swc_key : s3 key of the upright transformed morphology 
+        upright_png_key: s3 key of the morphology png
     """
 
     return {
         'upright_transform_dict': upright_transform,
         'upright_angle': upright_angle,
         "upright_swc_key": morphology_to_s3(
             bucket, f"{prefix}/upright_morphology.swc", upright_morphology
+        ),
+        "upright_png_key": morphology_png_to_s3(
+            bucket, f"{prefix}/upright_morphology.png", upright_morphology
         )
 
     }

pipeline/deploy/cloudformation/morphology_pipeline.yml

@@ -249,6 +249,29 @@ Resources:
         - FARGATE
       NetworkMode: awsvpc
 
+  PlotOriginalTask:
+    Type: AWS::ECS::TaskDefinition
+    Properties:
+      ContainerDefinitions:
+        - Image: !Ref ImageUri
+          Command: ["/bin/sh -c \"conda run -n py37 python /neuron_morphology/pipeline/commands/plot_original.py\""]
+          EntryPoint: ["sh", "-c"]
+          Name: plot-original-container
+          LogConfiguration:
+            LogDriver: awslogs
+            Options:
+              awslogs-group: !Sub "${AWS::StackName}-ecs"
+              awslogs-region: !Ref AWS::Region
+              awslogs-stream-prefix: plot-original-task
+      Cpu: 1024
+      ExecutionRoleArn: !GetAtt ContainerExecutionRole.Arn
+      TaskRoleArn: !GetAtt LandingTaskRole.Arn
+      Family: plot-original-task
+      Memory: 8192
+      RequiresCompatibilities:
+        - FARGATE
+      NetworkMode: awsvpc
+
   DepthFieldTask:
     Type: AWS::ECS::TaskDefinition
     Properties:
@@ -415,6 +438,49 @@ Resources:
                   }
                 },
                 "ResultPath": "$.landing",
+                "Next": "PlotOriginal"
+              },
+              "PlotOriginal": {
+                "Type": "Task",
+                "Resource": "arn:aws:states:::ecs:runTask.waitForTaskToken",
+                "HeartbeatSeconds": 2400,
+                "Parameters": {
+                  "Cluster": "${EcsArn}",
+                  "LaunchType": "FARGATE",
+                  "TaskDefinition": "${PlotOriginalEcsTask}",
+                  "Overrides": {
+                    "ContainerOverrides": [
+                      {
+                        "Name": "plot-original-container",
+                        "Environment": [
+                          {
+                            "Name": "TASK_TOKEN",
+                            "Value.$": "$$.Task.Token"
+                          },
+                          {
+                            "Name": "WORKING_BUCKET",
+                            "Value": "${WorkingBucket}"
+                          },
+                          {
+                            "Name": "RECONSTRUCTION_ID",
+                            "Value.$": "$.landing.reconstruction_id"
+                          },
+                          {
+                            "Name": "RUN_PREFIX",
+                            "Value.$": "$.landing.base_key"
+                          }
+                        ]
+                      }
+                    ]
+                  },
+                  "NetworkConfiguration": {
+                    "AwsvpcConfiguration": {
+                      "Subnets": ["${PrivateSubnet}"],
+                      "SecurityGroups": ["${ContainerSecurityGroup}"]
+                    }
+                  }
+                },
+                "ResultPath": "$.plot_original",
                 "Next": "DepthField"
               },
               "DepthField": {
@@ -566,7 +632,6 @@ Resources:
                 ],
                 "Default": "TiltCorrection"
               },
-
               "TiltCorrection": {
                 "Type": "Task",
                 "Resource": "arn:aws:states:::ecs:runTask.waitForTaskToken",
@@ -626,6 +691,7 @@ Resources:
           ContainerSecurityGroup: !Ref ContainerSecurityGroup
           WorkingBucket: !Ref WorkingBucket
           LandingBucket: !Ref LandingBucket
+          PlotOriginalEcsTask: !Ref PlotOriginalTask
           DepthFieldEcsTask: !Ref DepthFieldTask
           ScaleCorrectionTask: !Ref ScaleCorrectionTask
           UprightTransformEcsTask: !Ref UprightTransformTask