diff --git a/Python_script/prototype.py b/Python_script/prototype.py
index 1cef1c89ec09b8a086deefcd3be4b49f9354d658..9f3f7827594587d64a9c2918e9f8a874d3a1eb51 100755
--- a/Python_script/prototype.py
+++ b/Python_script/prototype.py
@@ -268,7 +268,7 @@ class Measurements:
 
         return equilibrium_indicator
 
-    def plot_output(self, output_file):
+    def plot_output(self, output_file, show_plot_window):
         csv_data = pd.read_csv(output_file)
         fig, ax1 = plt.subplots(2, figsize=(12, 10))
         fig.suptitle("Graphical representation of chamber output", color="red")
@@ -310,7 +310,9 @@ class Measurements:
         ax1[1].legend(all_path_collections, labels, loc='lower right')
 
         fig.savefig(time.strftime("%Y_%m_%d-%H_%M_%S") + 'graph.pdf')
-        plt.show()
+
+        if show_plot_window:
+            plt.show()
 
     def calculate_temperature_stability(self, target_temp, readback_temp):
         return (target_temp-self.max_delta_temp <= float(readback_temp)) and \
@@ -337,6 +339,9 @@ if __name__ == '__main__':
     parser.add_argument("-s", "--standby",
                         action="store_true", default=False,
                         help="Change Chamber mode to standby at the end")
+    parser.add_argument("-p", "--plot",
+                        action="store_true", default=False,
+                        help="Show plot window at the end")
 
     # parser.add_argument("-q", "--quiet",
     #                     action="store_false", dest="verbose", default=True,
@@ -354,17 +359,11 @@ if __name__ == '__main__':
     with open('test_stand_parameter.json', 'r') as f:
         config_data = json.load(f)
 
-    # Move this to inside the chamber. But just putting it into perform_measurements is not good at the moment.
-    # That function needs refactoring first.
     mes = Measurements(args.chamber, args.vna, args.file, output, args.standby, config_data)
     try:
 
         mes.perform_measurements()
-        plot_graph = input('want to plot a graph? ans[Yes/no]?').lower()
-        if plot_graph == "yes":
-            mes.plot_output(output)
-        else:
-            print("graph is stored with output file")
+        mes.plot_output(output, args.plot)
 
     finally:
         mes.chamber.close()