File size: 5,163 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 b4123b8 dd1d7f5 91a7a12 b4123b8 91a7a12 dd1d7f5 91a7a12 dd1d7f5 91a7a12 dd1d7f5 91a7a12 b4123b8 91a7a12 b4123b8 91a7a12 b4123b8 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 |
"""
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']
soft_mask = self.segmentation_manager.segment_image_soft(composite)
pdata['mask'] = (soft_mask * 255.0).astype(np.uint8)
return plants
def _extract_features(self, plants: Dict[str, Any]) -> Dict[str, Any]:
"""Extract texture, vegetation, and morphology features."""
for key, pdata in plants.items():
# Texture: LBP, HOG, Lacunarity from pseudo-color
composite = pdata['composite']
mask = pdata.get('mask')
masked = self.mask_handler.apply_mask_to_image(composite, mask) if mask is not None else composite
gray = cv2.cvtColor(masked, cv2.COLOR_BGR2GRAY)
feats = self.texture_extractor.extract_all_texture_features(gray)
stats = self.texture_extractor.compute_texture_statistics(feats, mask)
pdata['texture_features'] = {'color': {'features': feats, 'statistics': stats}}
# Vegetation: NDVI, ARI, GNDVI
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
pdata['morphology_features'] = self.morphology_extractor.extract_morphology_features(composite, mask)
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", "ARI", "GNDVI"):
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 |