bplusplus 1.1.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

bplusplus/__init__.py

@@ -1,5 +1,7 @@
 from .collect import Group, collect
 from .train_validate import train, validate
 from .prepare import prepare
-from .yolov5detect.detect import run
-
+from .resnet.train import train_resnet
+from .resnet.test import test_resnet
+from .hierarchical.train import train_multitask
+from .hierarchical.test import test_multitask

bplusplus/collect.py

@@ -1,11 +1,13 @@
 import os
 import random
+import threading
 from enum import Enum
-from typing import Any, Optional
+from typing import Any, Dict, List, Optional
 
 import pygbif
 import requests
 import validators
+from tqdm import tqdm
 
 
 #this lists currently supported groupings, more can be added with proper testing
@@ -13,10 +15,28 @@ class Group(str, Enum):
     scientificName="scientificName"
 
 #TODO add back support for fetching from dataset (or csvs)
-def collect(group_by_key: Group, search_parameters: dict[str, Any], images_per_group: int, output_directory: str):
+def collect(group_by_key: Group, search_parameters: dict[str, Any], images_per_group: int, output_directory: str, num_threads: int):
 
     groups: list[str] = search_parameters[group_by_key.value]
 
+    # check if user wants to parallelize the process
+    if num_threads > 1: 
+        __threaded_collect(
+            images_per_group=images_per_group, 
+            output_directory=output_directory, 
+            num_threads=num_threads,
+            groups=groups)
+    else: 
+        __single_collect(
+            search_parameters=search_parameters, 
+            images_per_group=images_per_group, 
+            output_directory=output_directory, 
+            group_by_key=group_by_key,
+            groups=groups,
+        )
+
+def __single_collect(group_by_key: Group, search_parameters: dict[str, Any], images_per_group: int, output_directory: str, groups: list[str]):
+
     #TODO throw error if groups is not a str list
 
     __create_folders(
@@ -37,7 +57,7 @@ def collect(group_by_key: Group, search_parameters: dict[str, Any], images_per_g
         sampled_occurrences = random.sample(occurrences, min(images_per_group, len(occurrences)))
         
         print(f"Downloading {len(sampled_occurrences)} images into the {group} folder...")
-        for occurrence in sampled_occurrences:
+        for occurrence in tqdm(sampled_occurrences, desc=f"Downloading images for {group}", unit="image"):
             # image_url = occurrence.image_url.replace("original", "large") # hack to get max 1024px image
 
             __down_image(
@@ -49,6 +69,38 @@ def collect(group_by_key: Group, search_parameters: dict[str, Any], images_per_g
     
     print("Finished collecting images.")
 
+# threaded_collect: paralellize the collection of images
+def __threaded_collect(images_per_group: int, output_directory: str, num_threads: int, groups: list[str]):
+    # Handle edge case where num_threads is greater than number of groups
+    if num_threads >= len(groups):
+        num_threads = len(groups)
+
+    # Divide the species list into num_threads parts
+    chunk_size = len(groups) // num_threads
+    species_chunks = [
+        groups[i:i + chunk_size] for i in range(0, len(groups), chunk_size)
+    ]
+    
+    # Ensure we have exactly num_threads chunks (the last chunk might be larger if len(species_list) % num_threads != 0)
+    while len(species_chunks) < num_threads:
+        species_chunks.append([])
+    
+    threads = []
+    for i, chunk in enumerate(species_chunks):
+        thread = threading.Thread(
+            target=__collect_subset,
+            args=(chunk, images_per_group, output_directory, i)
+        )
+        threads.append(thread)
+        thread.start()
+    
+    # Wait for all threads to complete
+    for thread in threads:
+        thread.join()
+    
+    print("All collection threads have finished.")
+    
+
 def _fetch_occurrences(group_key: str, group_value: str, parameters: dict[str, Any], totalLimit: int) -> list[dict[str, Any]]:
     parameters[group_key] = group_value
     return __next_batch(
@@ -98,6 +150,23 @@ def __create_folders(names: list[str], directory: str):
         # Create a folder using the group name
         os.makedirs(folder_name, exist_ok=True)
 
+def __collect_subset(species_subset: List[str], images_per_group: int, output_directory: str, thread_id: int):
+    search_subset: Dict[str, Any] = {
+        "scientificName": species_subset
+    }
+    
+    print(f"Thread {thread_id} starting collection for {len(species_subset)} species.")
+    
+    __single_collect(
+        search_parameters=search_subset,
+        images_per_group=images_per_group,
+        output_directory=output_directory,
+        group_by_key=Group.scientificName,
+        groups=species_subset
+    )
+    
+    print(f"Thread {thread_id} finished collection.")
+