update code

cvvsu · Nov 24, 2021 · a2b94ed · a2b94ed
1 parent 9eb6af3
commit a2b94ed
Show file tree

Hide file tree

Showing 9 changed files with 727 additions and 377 deletions.
diff --git a/README.md b/README.md
@@ -9,7 +9,6 @@ If you have any question, please open an issue or directly send an email to us.
 - Upload the pre-trained model used in the manuscript.
 - ~~Use a notebook to illustrate how to use the code step-by-step.~~
 - Train the Mask R-CNN model on all NPF events from the four datasets and upload the trained model.
-- Provide a fancy GUI.
 
 ## Requirements
 Please install the packages in the `requirements.txt`.
@@ -21,23 +20,8 @@ We find the versions of `PyTorch` (>=1.7) and `Python` (>=3.6) do not really aff
 
 ## Usage
 
-1. Please change the parameters in the `options.py` file. Specific parameters:
-
-    - `im_size` : image size for processing. It is also possible to use the size 128*128. 
-    - `scores`  : if you want analyze a long-term dataset, you can set a small value such as 0.2 or 0.3 to save time and effort. If you are dealing with one-year dataset, it is better to check all the detected masks (**0.0**).
-
-2. Run the `main.py`. 
-
-    - Drawing the NPF images does not take too much time.
-    - However, the mask detection process is quite slow if you do not have a powerful GPU. [Google Colab](https://colab.research.google.com/notebooks/intro.ipynb?utm_source=scs-index#) may provide a free GPU, which is a possible choice. After detecting the masks, you can download the masks to your local device for visualization.
-    - You need to select the masks yourself (more than one detected masks for some days).
-    - Once the masks are obtained, the GRs, start times, and end times can be determined automatically.
-
-**NOTE**: make sure that your local device is powerful enough; otherwise, please modify the function [save_SE_GR](https://github.com/cvvsu/maskNPF/blob/a959edf04f794d70e7ef8979494e8f36e317326e/model.py#L285) in the `model.py` to calculate GRs one-by-one (do not use the `multiprocessing` package).
-
-Or use the command line:
 ```
-$ python3 main.py --station hyytiala --im_size 256 --scores 0.0 --vmax 1e4
+$ python3 main.py --station hyytiala --im_size 256 --scores 0.0 --vmax 1e4 --model_name maskrcnn.pth
 ```
 
 Check the `demo.ipynb` file for more information.
@@ -57,17 +41,9 @@ We currently use a simple GUI (based on the `tk` package) to select the detected
 ## GRs, start times, and end times
 
 Our code can calculate GRs automatically for the size ranges: 
-
-[3-10 nm](https://github.com/cvvsu/maskNPF/blob/a959edf04f794d70e7ef8979494e8f36e317326e/model.py#L249)
-
-[10-25 nm](https://github.com/cvvsu/maskNPF/blob/a959edf04f794d70e7ef8979494e8f36e317326e/model.py#L250)
-
-[3-25 nm](https://github.com/cvvsu/maskNPF/blob/a959edf04f794d70e7ef8979494e8f36e317326e/model.py#L253)
-
-[size ranges determined by the detected masks](https://github.com/cvvsu/maskNPF/blob/a7c188f64e8329e0ae50ec23936158e4c89e07b0/model.py#L254)
-
-You can change the size range to calculate other GRs, taking 3-7 nm as an example. The parameters are changed in the fuction [get_GRs](https://github.com/cvvsu/maskNPF/blob/a959edf04f794d70e7ef8979494e8f36e317326e/model.py#L239) in the `model.py`. Once other size ranges are added, please also change the [save_dict](https://github.com/cvvsu/maskNPF/blob/a959edf04f794d70e7ef8979494e8f36e317326e/model.py#L274) at the same time in the [get_SE_GR](https://github.com/cvvsu/maskNPF/blob/a959edf04f794d70e7ef8979494e8f36e317326e/model.py#L257) function.
-
+- 3 - 10 nm
+- 10 - 25 nm
+- 3 - 25 nm
 
 
 ## Citation

diff --git a/main.py b/main.py
@@ -15,20 +15,20 @@
     #############################
     # draw NPF images
     #############################
-    NPF.draw_one_day_images()
-    NPF.draw_two_day_images()
+    #NPF.draw_one_day_images()
+    #NPF.draw_two_day_images()
 
 
     #############################
     # Get masks for NPF images
     #############################
-    NPF.detect_one_day_masks()
-    NPF.detect_two_day_masks()
+    # NPF.detect_one_day_masks()
+    # NPF.detect_two_day_masks()
 
     #############################
     # Visualize and select masks
     #############################
-    NPF.visualize_masks()
+    # NPF.visualize_masks()
 
     #####################################
     # Save the start-end times and GRs

diff --git a/model.py b/model.py
@@ -3,6 +3,7 @@
 import pandas as pd
 from multiprocessing import Pool
 from PIL import Image
+from tqdm import tqdm
 
 
 from matplotlib.figure import Figure
@@ -14,10 +15,10 @@
 
 import torch
 from torchvision import transforms
-from utils.utils import get_next_day, mkdirs, psd2im
-from utils.networks import get_instance_segmentation_model
-from utils.utils import reshape_mask
-from utils.fitting import get_GR, get_SE
+from utils import get_next_day, mkdirs, psd2im
+from utils import get_instance_segmentation_model
+from utils import reshape_mask
+from utils import get_GR, get_SE
 
 
 class NPFDetection(object):
@@ -29,7 +30,7 @@ def __init__(self, opt):
         self.cpu_count = os.cpu_count() // 2 + 1
         self.dataroot = os.path.join(opt.dataroot, opt.station)
         self.station = opt.station
-        self.vmax = opt.vmax
+        self.vmax = None if opt.dynamic_vmax else opt.vmax
         self.tm_res = opt.time_res
         self.df = pd.read_csv(os.path.join(self.dataroot, self.station+'.csv'), parse_dates=[0], index_col=0)         
         self.days = sorted(np.unique(self.df.index.date.astype(str)).tolist())
@@ -42,17 +43,26 @@ def draw_one_day_images(self):
         self.savefp = os.path.join(self.dataroot, 'images', 'one_day')
         mkdirs(self.savefp)
         self.dimg = 1
-
-        with Pool(self.cpu_count) as p:
-            p.map(self.draw_image, self.days)
+
+        if self.cpu_count >= 8:
+            with Pool(self.cpu_count) as p:
+                p.map(self.draw_image, self.days)
+        else:
+            for day in tqdm(self.days):
+                self.draw_image(day)
 
     def draw_two_day_images(self):
         """Draw NPF images with two-day unit"""
         self.savefp = os.path.join(self.dataroot, 'images', 'two_day')
         mkdirs(self.savefp)
         self.dimg = 2
-        with Pool(self.cpu_count) as p:
-            p.map(self.draw_image, self.days)
+
+        if self.cpu_count >= 8:
+            with Pool(self.cpu_count) as p:
+                p.map(self.draw_image, self.days)
+        else:
+            for day in tqdm(self.days):
+                self.draw_image(day)
 
     def draw_image(self, day):
         """Draw an NPF image"""
@@ -236,65 +246,50 @@ def save_mask(self, btn):
             mkdirs(savefp)
             np.save(os.path.join(savefp, f'{self.key}.npy'), self.masks_twoday[self.key][idx])
 
-    def get_GRs(self, df, mask):
-        mask = reshape_mask(mask, df.shape)
-        pos = np.where(mask)
-        ymin = np.min(pos[1])
-        ymax = np.max(pos[1])
-        dps = [float(dp) for dp in df.columns]
-        dp_start = dps[ymin]
-        dp_end = dps[ymax]
-        s_tm, e_tm = get_SE(df, mask)
-
-        gr_3_10, _, _ = get_GR(df, mask, 2.75, 10, savefp=None, tm_res=self.tm_res, vmax=self.vmax)
-        gr_10_25, _, _ = get_GR(df, mask, 10, 25.2, savefp=None, tm_res=self.tm_res, vmax=self.vmax)
-        savefp = os.path.join(self.dataroot, 'GR/3_25')
-        mkdirs(savefp)
-        gr_3_25, _, _ = get_GR(df, mask, 2.75, 25.2, savefp=savefp, tm_res=self.tm_res, vmax=self.vmax)
-        gr_mask, _, _ = get_GR(df, mask, dp_start, dp_end, savefp=None, tm_res=self.tm_res, vmax=self.vmax)
-        return s_tm, e_tm, gr_3_10, gr_10_25, gr_3_25, gr_mask
-
     def get_SE_GR(self, day):
         df = self.df.loc[day]
         mask = np.load(os.path.join(self.dataroot, 'masks/one_day', day+'.npy'), allow_pickle=True)
+        mask = reshape_mask(mask, df.shape)
         try:
-            s_tm, e_tm, gr_3_10, gr_10_25, gr_3_25, gr_mask = self.get_GRs(df, mask)
+            st, et = get_SE(df, mask)
+            gr_dict = get_GR(df, mask, self.tm_res, savefp=self.savefp, vmax=self.vmax)
         except:
-            print(day)
+            # print(day)
             return
 
         try:
             mask_ = np.load(os.path.join(self.dataroot, 'masks/two_day', day+'.npy'), allow_pickle=True)
             df_ = self.df.loc[day:get_next_day(day)]
             mask_ = reshape_mask(mask_, df_.shape)
-            st, et = get_SE(df_, mask_)
+            st_two, et_two = get_SE(df_, mask_)
         except:
-            st, et = s_tm, e_tm
-        #return s_tm, e_tm, st, et, gr_3_10, gr_10_25, gr_3_25
-        save_dict = {'date':[day],
-                     'GR (3-10)': [gr_3_10],
-                     'GR (10-25)': [gr_10_25],
-                     'GR (3-25)': [gr_3_25],
-                     'GR_mask': [gr_mask],
-                     'ST (one)': [s_tm],
-                     'ET (one)': [e_tm],
-                     'ST (two)': [st],
-                     'ET (two)': [et]}
+            st_two, et_two = st, et
+
+        save_dict = {**{
+            'date': [day],
+            'start_time_one': [st],
+            'end_time_one': [et],
+            'start_time_two': [st_two],
+            'end_time_two': [et_two]
+        }, **gr_dict}
         pd.DataFrame(save_dict).to_csv(os.path.join(self.savefp, f'{day}.csv'), index=False)
 
     def save_SE_GR(self):
-        files = glob.glob(os.path.join(self.dataroot, 'masks/one_day')+'/*.npy')
-        daysh = [file.split(os.sep)[-1].split('.')[0] for file in files]
-
-        grs = glob.glob(os.path.join(self.dataroot, 'GR')+'/*.csv')
-        days_1 = [gr.split(os.sep)[-1].split('.')[0] for gr in grs]
-
-        days = [item for item in daysh if item not in days_1]
-
+        r"""
+        obtain and save the start time, end time and the growth rates.
+        """
+        files = sorted(glob.glob(os.path.join(self.dataroot, 'masks/one_day')+'/*.npy'))
+        days = [file.split(os.sep)[-1].split('.')[0] for file in files]
+        print(f'Calculating growth rates for {len(days)} days.')
+
         self.savefp = os.path.join(self.dataroot, 'GR')
         mkdirs(self.savefp)
-        x_len = len(days) // 10
-        for i in range(11):
-            days_ =  days[i*x_len:(i+1)*x_len]
+
+        if self.cpu_count >= 8:
             with Pool(self.cpu_count) as p:
-                p.map(self.get_SE_GR, days_)
+                p.map(self.get_SE_GR, days)
+        else:
+            for day in tqdm(days):
+                self.get_SE_GR(day)
+
+
diff --git a/options.py b/options.py
@@ -24,10 +24,10 @@ def get_args():
     parser.add_argument('--im_size', default=256, type=int, help='image size for training. Square sizes are used but rectangle are also possible')
     parser.add_argument('--scores', default=0.00, type=float, help='threshold for objectiveness scores')
     parser.add_argument('--vmax', default=1e4, type=float, help='value scales for drawing')
+    parser.add_argument('--dynamic_vmax', action='store_true', help='utilize the dynamic surface plot if specified')
     parser.add_argument('--time_res', default=10, type=float, help='the time resolution of measurements. Default is 10 minutes')
     parser.add_argument('--mask_thres', default=0.5, type=float, help='threshold to binarize an mask')
     parser.add_argument('--ftsize', default=16, type=int, help='font size for elements in plots')
-    #parser.add_argument('--twoday', action='store_true', help='draw two-day NPF images')
     args = parser.parse_args()
 
     return args, get_message(parser, args)
diff --git a/requirements.txt b/requirements.txt
@@ -1,11 +1,11 @@
-astropy==4.2.1
-matplotlib==3.4.2
-pandas==1.2.4
-numpy==1.21.1
-torch==1.7.1
-tqdm==4.61.0
-seaborn==0.11.1
-scikit_image==0.18.1
-torchvision==0.8.2
-Pillow==8.3.1
-scikit_learn==0.24.2
+astropy>=4.2.1
+matplotlib>=3.4.2
+pandas>==1.2.4
+numpy>=1.21.1
+torch>=1.7.1
+tqdm>=4.61.0
+seaborn>=0.11.1
+scikit_image>=0.18.1
+torchvision>=0.8.2
+Pillow>=8.3.1
+scikit_learn>=0.24.2
diff --git a/utils/__init__.py b/utils/__init__.py
@@ -0,0 +1,4 @@
+from .networks import get_instance_segmentation_model
+from .visualize import psd2im, draw_subplots
+from .utils import mkdirs, get_next_day, reshape_mask, convert_matlab_time, time2num, num2time
+from .fitting import get_SE, get_GR, get_GR_old