feat: Add ViT models, HSI Hang2020, and prepare for DeepForest integration

BREAKING CHANGES:
- Default RGB image size changed from 128x128 to 224x224
- Default RGB normalization changed from 0_1 to imagenet
  For backward compatibility, explicitly pass rgb_size=(128, 128) and rgb_norm_method='0_1'

New Features:
- Add Vision Transformer (ViT) support: vit_b_16, vit_b_32, vit_l_16, vit_l_32
- Implement Hang2020 dual-pathway attention architecture for HSI classification
- Add model_variant parameter to training script for architecture selection
- Add preliminary DeepForest CropModel compatibility methods (WIP):
  * normalize() method for transforms
  * label_dict persistence in checkpoints
  * set_label_dict() and get_label_dict() helpers
- Add HuggingFace upload script (experimental, needs further testing)
- Add multi-output training support with auxiliary losses (Hang2020)

Improvements:
- Better experiment naming to prevent collisions in SLURM array jobs
- Enhanced test logging with detailed statistics
- Add rgb_size and rgb_norm_method CLI arguments for flexibility
- Update README with project roadmap

Note: Full DeepForest CropModel integration and HuggingFace loading
are still in progress and may require additional work.

Files changed: 10 files
- Added: scripts/upload_to_huggingface.py, sample_plots/test_PSMEM_douglas_fir.png
- Modified: train.py, rgb_models.py, hsi_models.py, lightning_modules.py,
            dataset.py, datamodule.py, README.md, visualization.ipynb

Author: Ritesh313

README.md

@@ -6,9 +6,9 @@ A comprehensive toolkit for multi-modal tree species classification using NEON e
 
 This repository aims to provide an end-to-end solution for tree species classification:
 
-- [x] **Dataset**: Ready-to-use multi-modal tree crown dataset with 167 species
+- [x] **Dataset**: Ready-to-use multi-modal tree crown dataset with 167 species. It's curated using the code in preprocessing directory in this repo.
 - [ ] **Data Processing**: Tools for downloading and processing raw NEON data products
-- [ ] **Classification Models**: Pre-trained models and training pipelines
+- [ ] **Classification Models**: Pre-trained models and training pipelines (Ongoing. ETA End of Feb 2026)
 - [ ] **DeepForest Integration**: Automated crown detection and classification workflow
 
 ## What's Available Now

examples/train.py

@@ -181,6 +181,12 @@ def main():
     parser.add_argument(
         "--model_type", type=str, default="simple", help="Model architecture type"
     )
+    parser.add_argument(
+        "--model_variant",
+        type=str,
+        default=None,
+        help="Model variant (e.g., 'vit_b_16', 'vit_l_16' for ViT models)",
+    )
     parser.add_argument(
         "--num_classes",
         type=int,
@@ -237,6 +243,23 @@ def main():
         action="store_true",
         help="Use WeightedRandomSampler for balanced class sampling (recommended for imbalanced datasets)",
     )
+    
+    # Image size arguments
+    parser.add_argument(
+        "--rgb_size",
+        type=int,
+        default=224,
+        help="RGB image size (single value for square images, e.g., 224 for 224x224). Default matches ImageNet pretraining.",
+    )
+    
+    # Normalization arguments
+    parser.add_argument(
+        "--rgb_norm_method",
+        type=str,
+        default="imagenet",
+        choices=["none", "0_1", "imagenet"],
+        help="RGB normalization method: 'imagenet' (recommended for pretrained models), '0_1' (simple [0,1] range), 'none'",
+    )
 
     # Reproducibility arguments
     parser.add_argument(
@@ -293,8 +316,10 @@ def main():
     worker_init_fn.base_seed = args.seed
 
     # Set up experiment name (auto-generate)
+    # Include model_variant and taxonomic_level to avoid collisions in array jobs
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    experiment_name = f"{args.modality}_{args.model_type}_{args.batch_size}_{timestamp}"
+    model_name = args.model_variant if args.model_variant else args.model_type
+    experiment_name = f"{args.modality}_{model_name}_{args.taxonomic_level}_{timestamp}"
 
     # Set up output directory with dynamic naming within provided path
     if args.output_dir is None:
@@ -324,6 +349,8 @@ def main():
         external_test_csv_path=args.external_test_csv,  # External test support
         external_test_hdf5_path=args.external_test_hdf5,  # External test support
         modalities=[args.modality],
+        rgb_size=(args.rgb_size, args.rgb_size),  # Image size for RGB
+        rgb_norm_method=args.rgb_norm_method,  # Normalization for RGB (imagenet for pretrained models)
         taxonomic_level=args.taxonomic_level,  # Species or genus level
         use_balanced_sampler=args.use_balanced_sampler,  # Balanced sampling
         split_method=args.split_method,
@@ -381,6 +408,11 @@ def main():
 
     # Create classifier based on modality
     if args.modality == "rgb":
+        # Prepare model kwargs
+        model_kwargs = {}
+        if args.model_variant is not None:
+            model_kwargs["model_variant"] = args.model_variant
+        
         classifier = RGBClassifier(
             model_type=args.model_type,
             num_classes=args.num_classes,
@@ -389,6 +421,8 @@ def main():
             scheduler=args.scheduler,
             weight_decay=args.weight_decay,
             log_images=True,  # Enable image logging for RGB
+            idx_to_label=datamodule.full_dataset.idx_to_label,  # For DeepForest CropModel compatibility
+            **model_kwargs,  # Pass model variant for ViT and other models
         )
     elif args.modality == "hsi":
         classifier = HSIClassifier(

neon_tree_classification/core/datamodule.py

@@ -60,13 +60,13 @@ def __init__(
         species_filter: Optional[List[str]] = None,
         site_filter: Optional[List[str]] = None,
         year_filter: Optional[List[int]] = None,
-        rgb_size: Tuple[int, int] = (128, 128),
+        rgb_size: Tuple[int, int] = (224, 224),  # Matches ImageNet pretraining
         hsi_size: Tuple[int, int] = (12, 12),
         lidar_size: Tuple[int, int] = (12, 12),
         rgb_resize_mode: str = "nearest",
         hsi_resize_mode: str = "nearest",
         lidar_resize_mode: str = "nearest",
-        rgb_norm_method: str = "0_1",
+        rgb_norm_method: str = "imagenet",  # ImageNet normalization for pretrained models
         hsi_norm_method: str = "per_sample",
         lidar_norm_method: str = "height",
         custom_transforms: Optional[Dict[str, Callable]] = None,

neon_tree_classification/core/dataset.py

@@ -44,15 +44,15 @@ def __init__(
         site_filter: Optional[List[str]] = None,
         year_filter: Optional[List[int]] = None,
         # Target sizes for training (required for consistent batching)
-        rgb_size: Tuple[int, int] = (128, 128),
+        rgb_size: Tuple[int, int] = (224, 224),  # Matches ImageNet pretraining
         hsi_size: Tuple[int, int] = (12, 12),
         lidar_size: Tuple[int, int] = (12, 12),
         # Resize methods (optimized for speed)
         rgb_resize_mode: str = "nearest",  # Fastest for RGB images
         hsi_resize_mode: str = "nearest",  # Changed to nearest for speed
         lidar_resize_mode: str = "nearest",  # Changed to nearest for speed
         # Normalization methods (performance-first defaults)
-        rgb_norm_method: str = "0_1",  # Simple division, fastest
+        rgb_norm_method: str = "imagenet",  # ImageNet normalization for pretrained models
         hsi_norm_method: str = "per_sample",  # Per-sample z-score, faster than per_pixel
         lidar_norm_method: str = "height",  # Simple max scaling, fastest
         # Custom transforms (optional, per-modality)