fix: add torchvision dependency and apply black formatting

Author: Ritesh313

examples/train.py

@@ -243,15 +243,15 @@ 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",
@@ -412,7 +412,7 @@ def main():
         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,

neon_tree_classification/core/datamodule.py

@@ -556,12 +556,12 @@ def train_dataloader(self) -> DataLoader:
         # Compute sampler if balanced sampling is enabled
         sampler = None
         shuffle = True
-        
+
         if self.use_balanced_sampler:
             print("⚖️  Using WeightedRandomSampler for balanced class sampling")
             sampler = self._create_weighted_sampler()
             shuffle = False  # Can't use shuffle with sampler
-        
+
         return DataLoader(
             self.train_dataset,
             batch_size=self.batch_size,
@@ -612,10 +612,10 @@ def test_dataloader(self) -> Optional[DataLoader]:
     def _create_weighted_sampler(self) -> WeightedRandomSampler:
         """
         Create WeightedRandomSampler for balanced class sampling.
-        
+
         Computes sample weights inversely proportional to class frequency,
         so rare classes are sampled more often and common classes less often.
-        
+
         Returns:
             WeightedRandomSampler for training dataset
         """
@@ -640,29 +640,30 @@ def _create_weighted_sampler(self) -> WeightedRandomSampler:
 
         # Count class frequencies
         class_counts = sample_labels.value_counts().to_dict()
-        
+
         # Compute weight for each class (inverse frequency)
         num_samples = len(sample_labels)
         class_weights = {
             cls: num_samples / count for cls, count in class_counts.items()
         }
-        
+
         # Assign weight to each sample based on its class
         sample_weights = [class_weights[label] for label in sample_labels]
         sample_weights = torch.DoubleTensor(sample_weights)
-        
+
         # Create sampler
         sampler = WeightedRandomSampler(
             weights=sample_weights,
             num_samples=len(sample_weights),
-            replacement=True  # Sample with replacement to oversample rare classes
+            replacement=True,  # Sample with replacement to oversample rare classes
         )
-        
+
         print(f"   Created sampler for {len(sample_weights)} samples")
-        print(f"   Sample weight range: {sample_weights.min():.3f} - {sample_weights.max():.3f}")
-        
-        return sampler
+        print(
+            f"   Sample weight range: {sample_weights.min():.3f} - {sample_weights.max():.3f}"
+        )
 
+        return sampler
 
     def get_class_weights(self) -> torch.Tensor:
         """
@@ -723,32 +724,32 @@ def get_class_weights(self) -> torch.Tensor:
     def _create_genus_label_mapping(self) -> Dict[str, int]:
         """
         Create genus-level label mapping from species names in the CSV.
-        
+
         Extracts genus (first word) from species_name column.
-        
+
         Returns:
             Dictionary mapping genus name to integer index
         """
         import warnings
-        
+
         # Load CSV to extract species names
         df = pd.read_csv(self.csv_path)
-        
+
         # Apply any filters that were specified
         if self.dataset_params.get("species_filter"):
             df = df[df["species"].isin(self.dataset_params["species_filter"])]
         if self.dataset_params.get("site_filter"):
             df = df[df["site"].isin(self.dataset_params["site_filter"])]
         if self.dataset_params.get("year_filter"):
             df = df[df["year"].isin(self.dataset_params["year_filter"])]
-        
+
         # Extract genus from species_name (first word)
         df["genus"] = df["species_name"].apply(lambda x: str(x).split()[0])
-        
+
         # Get unique genera and create mapping
         unique_genera = sorted(df["genus"].unique())
         label_to_idx = {genus: idx for idx, genus in enumerate(unique_genera)}
-        
+
         # Validate genus names and warn about edge cases
         non_alpha_genera = [g for g in unique_genera if not g.isalpha()]
         if non_alpha_genera:
@@ -758,7 +759,7 @@ def _create_genus_label_mapping(self) -> Dict[str, int]:
                 f"Run 'python processing/misc/inspect_labels.py' to review. "
                 f"To exclude, use: species_filter=[...]"
             )
-        
+
         # Check for known family names
         known_families = {"Pinaceae", "Rosaceae", "Fabaceae", "Asteraceae"}
         found_families = set(unique_genera) & known_families
@@ -769,7 +770,7 @@ def _create_genus_label_mapping(self) -> Dict[str, int]:
                 f"These represent unidentified species within that family. "
                 f"See docs/taxonomic_levels.md for more information."
             )
-        
+
         return label_to_idx
 
     def get_dataset_info(self) -> Dict[str, Any]:

neon_tree_classification/core/dataset.py

@@ -250,26 +250,32 @@ def _validate_species_consistency(self) -> None:
         # If the first mapping key is a species code (all uppercase, short), it's species-level
         # If it's a genus name (capitalized, longer), it's genus-level
         sample_label = next(iter(mapping_labels)) if mapping_labels else ""
-        is_genus_mapping = sample_label and sample_label[0].isupper() and sample_label[1:].islower()
-        
+        is_genus_mapping = (
+            sample_label and sample_label[0].isupper() and sample_label[1:].islower()
+        )
+
         if is_genus_mapping:
             # Genus-level mapping: validate that all species have extractable genus
             if "species_name" not in self.data.columns:
                 raise ValueError(
                     "Genus-level mapping detected but 'species_name' column not found in data. "
                     "Cannot extract genus from species names."
                 )
-            
+
             # Extract genera from species names and check they're all in mapping
-            data_genera = set(self.data["species_name"].apply(lambda x: str(x).split()[0]).unique())
+            data_genera = set(
+                self.data["species_name"].apply(lambda x: str(x).split()[0]).unique()
+            )
             missing_genera = data_genera - mapping_labels
             if missing_genera:
                 raise ValueError(
                     f"Genera extracted from dataset not found in external label mapping: {sorted(missing_genera)}. "
                     f"External mapping has: {sorted(mapping_labels)}"
                 )
-            
-            print(f"✓ Genus-level validation passed: All {len(data_genera)} genera found in mapping")
+
+            print(
+                f"✓ Genus-level validation passed: All {len(data_genera)} genera found in mapping"
+            )
         else:
             # Species-level mapping: check species codes
             missing_in_mapping = data_species - mapping_labels

neon_tree_classification/inference/__init__.py

@@ -24,11 +24,11 @@
 from .utils import load_label_mapping, format_predictions
 
 __all__ = [
-    'TreeClassifier',
-    'preprocess_image',
-    'prepare_tensor',
-    'load_label_mapping',
-    'format_predictions',
+    "TreeClassifier",
+    "preprocess_image",
+    "prepare_tensor",
+    "load_label_mapping",
+    "format_predictions",
 ]
 
-__version__ = '1.0.0'
+__version__ = "1.0.0"