From 234ecc91ef98ebe813aff23c538fb799732bb741 Mon Sep 17 00:00:00 2001
From: GenevieveBuckley <30920819+GenevieveBuckley@users.noreply.github.com>
Date: Wed, 29 Jan 2020 18:21:07 +1100
Subject: [PATCH] Support command line use as described in the README

---
 segmentify/main.py | 16 +++++++++++-----
 setup.py           |  1 +
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/segmentify/main.py b/segmentify/main.py
index d1d875c..376bbf7 100644
--- a/segmentify/main.py
+++ b/segmentify/main.py
@@ -20,13 +20,19 @@ def main(args):
         imgs = np.array(imgs)
 
     with gui_qt():
-        viewer = Viewer(imgs, heatmap=args.heatmap)
-
-
-if __name__ == "__main__":
+        try:
+            args.heatmap
+        except AttributeError:
+            viewer = Viewer(imgs)
+        else:
+            viewer = Viewer(imgs, heatmap=args.heatmap)
+
+def main_cli():
     # parser
     parser = argparse.ArgumentParser()
     parser.add_argument("images", nargs="*", type=str, help="Image to view and segment.")
     args = parser.parse_args()
-
     main(args)
+
+if __name__ == "__main__":
+    main_cli()
diff --git a/setup.py b/setup.py
index 6fb822a..be53508 100644
--- a/setup.py
+++ b/setup.py
@@ -43,4 +43,5 @@ def parse_requirements_file(filename):
         'Operating System :: MacOS',
     ],
     install_requires=INSTALL_REQUIRES,
+    entry_points={'console_scripts': ['segmentify=segmentify.main:main_cli']},
 )