hcast model in pipeline

Author: bw4sz

boem_conf/boem_config.yaml

@@ -1,6 +1,6 @@
 defaults:
   - server: serenity
-  - classification_model: finetune.yaml # or hierarchical.yaml
+  - classification_model: finetune.yaml
   - annotation: label_studio
 
 hydra:

boem_conf/classification_model/hierarchical.yaml

@@ -1,7 +1,10 @@
 backend: hierarchical
+enabled: False
 force_train: False
 # Optional: explicitly set checkpoint path; if empty we auto-discover under tamu_hcast/
 checkpoint: /home/b.weinstein/BOEM/tamu_hcast/output/usgs_hcast_300_b256/best_checkpoint.pth
+# Optional path to CSV with species/genus/family labels
+label_csv: null
 # Batch size for classifying crops after detection
 batch_size: 64
 workers: 4

pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   # H-CAST requirements
   "timm",
   "opencv-contrib-python",
-  #"deepforest",
+  "deepforest<2.0",
   # DGL requires CUDA-specific wheels - will be resolved from find-links configured below
   "dgl",
   "packaging",
@@ -42,6 +42,7 @@ find-links = [
 ]
 # Exclude
 override-dependencies = [
-  "opencv-python; python_version < '0'"
+  "opencv-python; python_version < '0'",
+  "opencv-python-headless; python_version < '0'",
 ]
 

src/active_learning.py

@@ -1,5 +1,6 @@
 import random
 from src import detection
+from src import hierarchical
 
 def human_review(predictions, min_detection_score=0.6, min_classification_score=0.5, confident_threshold=0.5):
     """
@@ -27,7 +28,7 @@ def human_review(predictions, min_detection_score=0.6, min_classification_score=
 
     return confident_predictions, uncertain_predictions
 
-def generate_pool_predictions(pool, patch_size=512, patch_overlap=0.1, min_score=0.1, model=None, batch_size=16, pool_limit=1000, crop_model=None):
+def generate_pool_predictions(pool, patch_size=512, patch_overlap=0.1, min_score=0.1, model=None, batch_size=16, pool_limit=1000, crop_model=None, hcast_model=None, image_dir=None, hcast_batch_size=None, hcast_workers=None):
     """
     Generate predictions for the flight pool.
     
@@ -41,9 +42,13 @@ def generate_pool_predictions(pool, patch_size=512, patch_overlap=0.1, min_score
         comet_logger (CometLogger, optional): A CometLogger object. Defaults to None.
         crop_model (bool, optional): A deepforest.model.CropModel object. Defaults to None.
         pool_limit (int, optional): The maximum number of images to consider. Defaults to 1000.
+        hcast_model (optional): H-CAST hierarchical model wrapper. Defaults to None.
+        image_dir (str, optional): Root directory where images are located. Required if hcast_model is provided.
+        hcast_batch_size (int, optional): Batch size for H-CAST classification. Defaults to 64.
+        hcast_workers (int, optional): Number of workers for H-CAST DataLoader. Defaults to 4.
     
     Returns:
-        pd.DataFrame: A DataFrame of predictions.
+        pd.DataFrame: A DataFrame of predictions with both cropmodel and hcast columns (if hcast_model provided).
     """
 
     #subsample
@@ -63,6 +68,18 @@ def generate_pool_predictions(pool, patch_size=512, patch_overlap=0.1, min_score
 
     preannotations = preannotations[preannotations["score"] >= min_score]
 
+    # Apply hierarchical classification if hcast_model is provided
+    if hcast_model is not None:
+        if image_dir is None:
+            raise ValueError("image_dir is required when hcast_model is provided")
+        preannotations = hierarchical.classify_dataframe(
+            predictions=preannotations,
+            image_dir=image_dir,
+            model=hcast_model,
+            batch_size=hcast_batch_size,
+            num_workers=hcast_workers,
+        )
+
     return preannotations
 
 def select_images(preannotations, strategy, n=10, target_labels=None, min_score=0.3):

src/hierarchical.py

@@ -2,15 +2,15 @@
 from typing import Optional, Tuple, List, Dict, Callable
 
 import torch
-from torch.utils.data import Dataset, DataLoader
-from PIL import Image
-from src.hcast.cast_models import cast_deit_hier  
 import pandas as pd
 import numpy as np
+from PIL import Image
 import cv2
+
+from torch.utils.data import Dataset, DataLoader
 from torchvision import transforms
 from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
-
+from timm.models import create_model
 
 def _infer_head_sizes_from_checkpoint(ckpt: Dict[str, torch.Tensor]) -> Tuple[int, Optional[int], Optional[int]]:
     species = None
@@ -167,9 +167,6 @@ def load_hcast_model(
         model_state_dict = checkpoint["state_dict"]
     else:
         model_state_dict = checkpoint
-
-    # Get the training arguments if available
-    from timm.models import create_model
 
     if 'args' in checkpoint:
         args = checkpoint['args']

src/pipeline.py

@@ -139,7 +139,10 @@ def run(self):
         val_crop_image_dir = os.path.join(self.config.classification_model.val_crop_image_dir, self.comet_logger.experiment.id)
         os.makedirs(val_crop_image_dir, exist_ok=True)
 
-        if classification_backend == "deepforest":
+        # Always load the finetune/cropmodel (load both models when both configs are present)
+        trained_classification_model = None
+        # Load cropmodel if finetune config is available (which it should be by default)
+        if hasattr(self.config.classification_model, "checkpoint") and self.config.classification_model.checkpoint:
             # If there are no train annotations, turn off force training
             if all_training.xmin[all_training.xmin != 0].empty:
                 self.config.classification_model.force_train = False
@@ -170,12 +173,23 @@ def run(self):
                     workers=self.config.classification_model.workers,
                     comet_logger=self.comet_logger)
             else:
-                trained_classification_model = CropModel.load_from_checkpoint(self.config.classification_model.checkpoint )
-        else:
-            # Hierarchical backend (H-CAST). Load wrapper and classify crops post-detection.
-            repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
-            hcast_checkpoint = getattr(self.config.classification_model, "checkpoint", None)
-            hcast_wrapper = hierarchical.load_hcast_model(repo_root=repo_root, checkpoint_path=hcast_checkpoint)
+                trained_classification_model = CropModel.load_from_checkpoint(self.config.classification_model.checkpoint)
+
+        # Load hierarchical model (H-CAST) if enabled
+        hcast_model = None
+        if classification_backend == "hierarchical":
+            hcast_enabled = getattr(self.config.classification_model, "enabled", False)
+            if hcast_enabled:
+                hcast_checkpoint = getattr(self.config.classification_model, "checkpoint", None)
+                if hcast_checkpoint:
+                    label_csv = getattr(self.config.classification_model, "label_csv", None)
+                    hcast_model = hierarchical.load_hcast_model(
+                        checkpoint_path=hcast_checkpoint,
+                        label_csv=label_csv
+                    )
+                    print(f"Loaded H-CAST model from {hcast_checkpoint}")
+                else:
+                    print("H-CAST enabled but no checkpoint path provided, skipping hierarchical model")
 
         pool = glob.glob(os.path.join(self.config.image_dir, "*.jpg"))  # Get all images in the data directory
         pool = [image for image in pool if not image.endswith('.csv')]
@@ -189,6 +203,13 @@ def run(self):
             else:
                 pool = random.sample(pool, 10)
 
+        # Get hierarchical model parameters from config (hierarchical.yaml)
+        hcast_batch_size = None
+        hcast_workers = None
+        if classification_backend == "hierarchical" and getattr(self.config.classification_model, "enabled", False):
+            hcast_batch_size = self.config.classification_model.batch_size
+            hcast_workers = self.config.classification_model.workers
+
         flightline_predictions = generate_pool_predictions(
             pool=pool,
             pool_limit=self.config.active_learning.pool_limit,
@@ -198,6 +219,10 @@ def run(self):
             model=trained_detection_model,
             batch_size=self.config.predict.batch_size,
             crop_model=trained_classification_model,
+            hcast_model=hcast_model,
+            image_dir=self.config.image_dir if hcast_model is not None else None,
+            hcast_batch_size=hcast_batch_size,
+            hcast_workers=hcast_workers,
         )
 
         if flightline_predictions is None:
@@ -209,24 +234,26 @@ def run(self):
         if self.existing_validation is None:
             print("No validation annotations, skipping evaluation")       
         else:
-            evaluation_annotations = self.existing_validation.copy(deep=True)
-            evaluation_predictions = flightline_predictions[flightline_predictions.image_path.isin(self.existing_validation.image_path)]
-
-
-            label_dict = trained_classification_model.label_dict
-                
-            pipeline_monitor = PipelineEvaluation(
-                predictions=evaluation_predictions,
-                annotations=evaluation_annotations,
-                classification_label_dict=label_dict,
-                **self.config.pipeline_evaluation)
-
-            performance = pipeline_monitor.evaluate()
-            self.comet_logger.experiment.log_metrics(performance)
-
-            if pipeline_monitor.check_success():
-                print("Pipeline performance is satisfactory, exiting")
-                return None
+            if trained_classification_model is None:
+                print("No classification model available, skipping evaluation")
+            else:
+                evaluation_annotations = self.existing_validation.copy(deep=True)
+                evaluation_predictions = flightline_predictions[flightline_predictions.image_path.isin(self.existing_validation.image_path)]
+
+                label_dict = trained_classification_model.label_dict
+                    
+                pipeline_monitor = PipelineEvaluation(
+                    predictions=evaluation_predictions,
+                    annotations=evaluation_annotations,
+                    classification_label_dict=label_dict,
+                    **self.config.pipeline_evaluation)
+
+                performance = pipeline_monitor.evaluate()
+                self.comet_logger.experiment.log_metrics(performance)
+
+                if pipeline_monitor.check_success():
+                    print("Pipeline performance is satisfactory, exiting")
+                    return None
 
         test_preannotations = flightline_predictions[~flightline_predictions.image_path.isin(self.existing_images)]
         test_images_to_annotate, preannotations = select_images(