cleaned notebook (moved generate_synthetic_data to a separate script file)

doc/source/_static/example/load_tracking.py

+##########################################################################
+# Track Analyzer - Quantification and visualization of tracking data     #
+# Authors: Arthur Michaut                                                #
+# Copyright 2016-2019 Harvard Medical School and Brigham and             #
+#                          Women's Hospital                              #
+# Copyright 2019-2021 Institut Pasteur and CNRS–UMR3738                  #
+# See the COPYRIGHT file for details                                     #
+#                                                                        #
+# This file is part of Track Analyzer package.                           #
+#                                                                        #
+# Track Analyzer is free software: you can redistribute it and/or modify #
+# it under the terms of the GNU General Public License as published by   #
+# the Free Software Foundation, either version 3 of the License, or      #
+# (at your option) any later version.                                    #
+#                                                                        #
+# Track Analyzer is distributed in the hope that it will be useful,      #
+# but WITHOUT ANY WARRANTY; without even the implied warranty of         #
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the           #
+# GNU General Public License for more details .                          #
+#                                                                        #
+# You should have received a copy of the GNU General Public License      #
+# along with Track Analyzer (COPYING).                                   #
+# If not, see <https://www.gnu.org/licenses/>.                           #
+##########################################################################
+
 import os.path as osp
 import sys
 import argparse

track_analyzer/synthetic_data.py

@@ -24,12 +24,12 @@
 ##########################################################################
 
 import os.path as osp
-import shutil
 
 import matplotlib.pyplot as plt
+import matplotlib as mpl
 import numpy as np
 import pandas as pd
-from skimage import io
+from skimage.color import rgb2gray
 from skimage.util import img_as_ubyte
 import seaborn as sns
 import tifffile as tifff
@@ -42,6 +42,7 @@ color_list = [c['color'] for c in list(plt.rcParams['axes.prop_cycle'])] + sns.c
 plot_param = {'figsize': (5, 5), 'dpi': 300, 'color_list': color_list, 'format': '.png', 'despine': True, 'logx': False,
               'logy': False, 'invert_yaxis': True, 'export_data_pts': False}
 
+mpl.use('agg')
 
 def make_diff_traj(part_index=0, grid_size=[500, 500, 500], dim=3, tmax=10, periodic=True, noise_amp=10,
                    x0=[250, 250, 250], bias=[0, 0, 0]):
@@ -173,17 +174,16 @@ def make_attraction_node(part_num=100, grid_size=[500, 500, 500], dim=3, tmax=10
 
 def plot_synthetic_stack(df, outdir, dpi=300, grid_size=[500, 500, 500], tmax=10):
     """Plot synthetic data and save it as a grayscaled tiff stack"""
-    outdir_temp = osp.join(outdir, 'temp')
-    tpr.safe_mkdir(outdir_temp)
-
-    stack = np.zeros((tmax, grid_size[0], grid_size[1]), 'uint8')
 
+    stack = []  # stack to store images
     groups = df.groupby('frame')
-    # print
+    
+    # plot frames
     for i in range(tmax):
         group = groups.get_group(i).reset_index(drop=True)
-        fig = plt.figure(frameon=False)
-        fig.set_size_inches(grid_size[0] / dpi, grid_size[1] / dpi)
+
+        figsize = (grid_size[0] / dpi, grid_size[1] / dpi)
+        fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)
         ax = fig.add_axes([0, 0, 1, 1])
         for k in range(group.shape[0]):
             ax.scatter(group.loc[k, 'x'], group.loc[k, 'y'], s=10)
@@ -191,15 +191,17 @@ def plot_synthetic_stack(df, outdir, dpi=300, grid_size=[500, 500, 500], tmax=10
         ax.set_ylim(0, grid_size[1])
         ax.invert_yaxis()
         ax.axis('off')
-        fn = osp.join(outdir_temp, '{}.jpg'.format(i))
-        fig.savefig(fn, dpi=300)
+        
+        fig.canvas.draw()
+        fig_image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+        fig_image = fig_image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        fig_image = rgb2gray(fig_image)  # convert to grayscale
+        fig_image = img_as_ubyte(fig_image)  # convert to 8 bit
+        stack.append(fig_image)
+        plt.close(fig)
+
+    # save
+    stack = np.array(stack)
+    tifff.imsave(osp.join(outdir,'stack.tiff'), stack)
 
-    # add to stack
-    for i in range(tmax):
-        fn = osp.join(outdir_temp, '{}.jpg'.format(i))
-        im = io.imread(fn, as_gray=True)
-        stack[i] = img_as_ubyte(im)
 
-    out_fn = osp.join(outdir, 'stack.tiff')
-    tifff.imsave(out_fn, stack)
-    shutil.rmtree(outdir_temp)