Add progress callback to search_images

Author: jrrodri

abraia/training/dataset.py

@@ -85,7 +85,7 @@ def search_google(query):
         yield link
 
 
-def download(query, limit=100, save_output='dataset', verbose=True):
+def download(query, limit=100, save_output='dataset', verbose=True, callback=None):
     seen = set()
     download_count = 0
     os.makedirs(save_output, exist_ok=True)
@@ -102,6 +102,8 @@ def download(query, limit=100, save_output='dataset', verbose=True):
                         download_count += 1
                         if verbose:
                             print(f"[%] Downloaded Image #{download_count} from {link}")
+                        if callback:
+                            callback({'current': download_count, 'total': limit})
                     except Exception as e:
                         print(f"[!] Error getting {link}: {e}")
                 else:
@@ -112,9 +114,9 @@ def download(query, limit=100, save_output='dataset', verbose=True):
                 break
 
 
-def search_images(query, limit=100, save_output='dataset', verbose=True):
+def search_images(query, limit=100, save_output='dataset', verbose=True, callback=None):
     """Search and download images from Google and Bing."""
-    download(query, limit=limit, save_output=save_output,  verbose=verbose)
+    download(query, limit=limit, save_output=save_output, verbose=verbose, callback=callback)
     return list_dir(save_output)
 
 

abraia/training/ops.py

@@ -0,0 +1,55 @@
+import numpy as np
+
+
+def train_test_split(*arrays, test_size=None, train_size=None, random_state=None, shuffle=True, stratify=None):
+    """Split arrays or matrices into random train and test subsets"""
+    if random_state is not None:
+        np.random.seed(random_state)
+    n_samples = len(arrays[0])
+    if test_size is None and train_size is None:
+        test_size = 0.25
+    if train_size is not None:
+        n_train = int(train_size * n_samples) if isinstance(train_size, float) else train_size
+        n_test = n_samples - n_train
+    else:
+        n_test = int(test_size * n_samples) if isinstance(test_size, float) else test_size
+        n_train = n_samples - n_test
+    indices = np.arange(n_samples)
+    if shuffle:
+        if stratify is not None:
+            unique_classes, class_indices = np.unique(stratify, return_inverse=True)
+            train_indices, test_indices = [], []
+            for i in range(len(unique_classes)):
+                cls_indices = indices[class_indices == i]
+                np.random.shuffle(cls_indices)
+                n_cls_test = int(n_test * len(cls_indices) / n_samples)
+                test_indices.extend(cls_indices[:n_cls_test])
+                train_indices.extend(cls_indices[n_cls_test:])
+            indices = np.array(train_indices + test_indices)
+            n_train = len(train_indices)
+        else:
+            np.random.shuffle(indices)
+    res = []
+    for arr in arrays:
+        if isinstance(arr, list):
+            res.append([arr[i] for i in indices[:n_train]])
+            res.append([arr[i] for i in indices[n_train:]])
+        else:
+            res.append(arr[indices[:n_train]])
+            res.append(arr[indices[n_train:]])
+    return res
+
+
+def resample(*arrays, n_samples=None, random_state=None, replace=True):
+    """Resample arrays or matrices in a consistent way"""
+    if random_state is not None:
+        np.random.seed(random_state)
+    n_samples = n_samples or len(arrays[0])
+    indices = np.random.choice(len(arrays[0]), size=n_samples, replace=replace)
+    res = []
+    for arr in arrays:
+        if isinstance(arr, list):
+            res.append([arr[i] for i in indices])
+        else:
+            res.append(arr[indices])
+    return res[0] if len(res) == 1 else res

scripts/abraia

@@ -200,10 +200,9 @@ def create(project, query, anonymize=False, upscale=0):
     """Create or update a dataset."""
     from abraia.training import search_images, load_dataset
     dataset = load_dataset(project)
-    save_output = f"{project}/"
     files = []
     if query:
-        files = search_images(query, limit=100, save_output=save_output, verbose=True)
+        files = search_images(query, limit=100, save_output=f"{project}/", verbose=True)
     if os.path.exists(project):
         files = process_dataset(project, anonymize=anonymize, upscale_threshold=upscale)
     process_map(upload_file, files, itertools.repeat(project + '/'), desc="Uploading images")