refactor: Improve verbosity messages and progress bars in data processing functions

Author: chanshing

src/stepcount/models.py

@@ -493,14 +493,14 @@ def make_windows(data, window_sec, fn=None, return_index=False, verbose=True):
     """ Split data into windows """
 
     if verbose:
-        print("Defining windows...")
+        print("Defining segments...")
 
     if fn is None:
         def fn(x):
             return x
 
     X, T = [], []
-    for t, x in tqdm(data.resample(f"{window_sec}s", origin="start"), mininterval=5, disable=not verbose):
+    for t, x in data.resample(f"{window_sec}s", origin="start"):
         x = fn(x)
         X.append(x)
         T.append(t)
@@ -626,12 +626,10 @@ def smooth_mean_absolute_percentage_error(yt, yp, sample_weight=None):
 def batch_extract_features(X, sample_rate, to_numpy=True, n_jobs=1, verbose=False):
     """ Extract features for a list or array of windows """
 
-    if verbose:
-        print("Extracting features...")
 
     X_feats = Parallel(n_jobs=n_jobs)(
         delayed(features.extract_features)(x, sample_rate)
-        for x in tqdm(X, mininterval=5, disable=not verbose)
+        for x in tqdm(X, total=len(X), mininterval=5, disable=not verbose, bar_format='Extracting features: {percentage:3.0f}%|{bar}| [{elapsed}<{remaining}]')
     )
     X_feats = pd.DataFrame(X_feats)
 

src/stepcount/sslmodel.py

@@ -237,8 +237,6 @@ def predict(model, dataloader, device, output_logits=False):
     :rtype: (np.ndarray, np.ndarray, np.ndarray)
     """
 
-    if verbose:
-        print('Classifying windows...')
 
     predictions_list = []
     true_list = []
@@ -249,7 +247,7 @@ def predict(model, dataloader, device, output_logits=False):
         return np.array([]), np.array([]), np.array([])
 
     with torch.inference_mode():
-        for x, y, pid in tqdm(dataloader, mininterval=5, disable=not verbose):
+        for x, y, pid in tqdm(dataloader, total=len(dataloader), mininterval=5, disable=not verbose, bar_format='Classifying segments: {percentage:3.0f}%|{bar}| [{elapsed}<{remaining}]'):
             x = x.to(device, dtype=torch.float)
             logits = model(x)
             true_list.append(y)
@@ -314,7 +312,7 @@ def train(model, train_loader, val_loader, device, class_weights=None, weights_p
         model.train()
         train_losses = []
         train_acces = []
-        for x, y, _ in tqdm(train_loader, disable=not verbose):
+        for x, y, _ in tqdm(train_loader, total=len(train_loader), disable=not verbose, bar_format='Training: {percentage:3.0f}%|{bar}| [{elapsed}<{remaining}]'):
             x.requires_grad_(True)
             x = x.to(device, dtype=torch.float)
             true_y = y.to(device, dtype=torch.long)

src/stepcount/stepcount.py

@@ -153,8 +153,6 @@ def main():
 
     model.wd.device = args.pytorch_device
 
-    if verbose:
-        print("Running step counter...")
     Y, W, T_steps = model.predict_from_frame(data)
 
     # Save step counts