File size: 5,854 Bytes
b4123b8 91a7a12 b4123b8 91a7a12 b4123b8 dd1d7f5 b4123b8 dd1d7f5 b4123b8 91a7a12 b4123b8 dd1d7f5 91a7a12 dd1d7f5 b4123b8 91a7a12 dd1d7f5 b4123b8 dffab99 b4123b8 dd1d7f5 b4123b8 91a7a12 b4123b8 dd1d7f5 91a7a12 b4123b8 91a7a12 b4123b8 91a7a12 b4123b8 91a7a12 b4123b8 91a7a12 b4123b8 91a7a12 2c0bae7 91a7a12 2c0bae7 b4123b8 dd1d7f5 10bba96 b4123b8 dd1d7f5 c170961 10bba96 c170961 31ddfa7 c170961 91a7a12 31ddfa7 b4123b8 91a7a12 b4123b8 91a7a12 c170961 91a7a12 b4123b8 31ddfa7 91a7a12 b4123b8 91a7a12 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
"""
Minimal single-image pipeline for Hugging Face demo.
"""
import logging
from pathlib import Path
from typing import Dict, Any
import numpy as np
import cv2
from .config import Config
from .data import ImagePreprocessor, MaskHandler
from .features import TextureExtractor, VegetationIndexExtractor, MorphologyExtractor
from .output import OutputManager
from .segmentation import SegmentationManager
logger = logging.getLogger(__name__)
class SorghumPipeline:
"""Minimal pipeline for single-image processing."""
def __init__(self, config: Config):
"""Initialize pipeline."""
logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s')
self.config = config
self.config.validate()
# Initialize components with defaults
self.preprocessor = ImagePreprocessor()
self.mask_handler = MaskHandler()
self.texture_extractor = TextureExtractor()
self.vegetation_extractor = VegetationIndexExtractor()
self.morphology_extractor = MorphologyExtractor()
self.segmentation_manager = SegmentationManager(
model_name="briaai/RMBG-2.0",
device=self.config.get_device(),
trust_remote_code=True
)
self.output_manager = OutputManager(
output_folder=self.config.paths.output_folder,
settings=self.config.output
)
logger.info("Pipeline initialized")
def run(self, single_image_path: str) -> Dict[str, Any]:
"""Run pipeline on single image."""
logger.info("Processing single image...")
from PIL import Image
import time
start = time.perf_counter()
# Load image
img = Image.open(single_image_path)
plants = {
"demo": {
"raw_image": (img, Path(single_image_path).name),
"plant_name": "demo",
}
}
# Process: composite → segment → features → save
plants = self.preprocessor.create_composites(plants)
plants = self._segment(plants)
plants = self._extract_features(plants)
self.output_manager.create_output_directories()
for key, pdata in plants.items():
self.output_manager.save_plant_results(key, pdata)
elapsed = time.perf_counter() - start
logger.info(f"Completed in {elapsed:.2f}s")
return {"plants": plants, "timing": elapsed}
def _segment(self, plants: Dict[str, Any]) -> Dict[str, Any]:
"""Segment using BRIA."""
for key, pdata in plants.items():
composite = pdata['composite']
logger.info(f"Composite shape: {composite.shape}")
soft_mask = self.segmentation_manager.segment_image_soft(composite)
logger.info(f"Soft mask shape: {soft_mask.shape}")
mask_uint8 = (soft_mask * 255.0).astype(np.uint8)
logger.info(f"Mask uint8 shape: {mask_uint8.shape}")
pdata['mask'] = mask_uint8
return plants
def _extract_features(self, plants: Dict[str, Any]) -> Dict[str, Any]:
"""Extract features (NDVI only for now)."""
for key, pdata in plants.items():
composite = pdata['composite']
mask = pdata.get('mask')
# Texture: ONLY LBP on green band within mask
pdata['texture_features'] = {}
green_band = None
spectral = pdata.get('spectral_stack', {})
if 'green' in spectral:
green_band = spectral['green'].squeeze(-1).astype(np.float64)
if mask is not None:
valid = np.where(mask > 0, green_band, np.nan)
else:
valid = green_band
# normalize to uint8 for LBP
v = valid.copy()
v = np.nan_to_num(v, nan=np.nanmin(v))
m, M = np.min(v), np.max(v)
denom = (M - m) if (M - m) > 1e-6 else 1.0
gray8 = ((v - m) / denom * 255.0).astype(np.uint8)
lbp_map = self.texture_extractor.extract_lbp(gray8)
pdata['texture_features'] = {'green': {'features': {'lbp': lbp_map}}}
# Vegetation: NDVI, GNDVI, SAVI
spectral = pdata.get('spectral_stack', {})
if spectral and mask is not None:
pdata['vegetation_indices'] = self._compute_vegetation(spectral, mask)
else:
pdata['vegetation_indices'] = {}
# # Morphology: PlantCV size analysis (COMMENTED OUT)
# pdata['morphology_features'] = self.morphology_extractor.extract_morphology_features(composite, mask)
pdata['morphology_features'] = {}
return plants
def _compute_vegetation(self, spectral: Dict[str, np.ndarray], mask: np.ndarray) -> Dict[str, Any]:
"""Compute NDVI, ARI, GNDVI only."""
out = {}
for name in ("NDVI", "GNDVI", "SAVI"):
bands = self.vegetation_extractor.index_bands.get(name, [])
if not all(b in spectral for b in bands):
continue
arrays = [np.asarray(spectral[b].squeeze(-1), dtype=np.float64) for b in bands]
values = self.vegetation_extractor.index_formulas[name](*arrays).astype(np.float64)
binary_mask = (mask > 0)
masked_values = np.where(binary_mask, values, np.nan)
valid = masked_values[~np.isnan(masked_values)]
stats = {
'mean': float(np.mean(valid)) if valid.size else 0.0,
'std': float(np.std(valid)) if valid.size else 0.0,
}
out[name] = {'values': masked_values, 'statistics': stats}
return out |