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']},
 )