| import yaml |
| import cv2 |
| import os |
| import torch |
| from .cnn_model import load_model |
| from .sam_utils import get_sam_generator |
| from .greenery_estimator import GreeneryEstimator |
|
|
| class EcoPulsePipeline: |
| """ |
| Orchestration layer for the EcoPulse satellite analysis system. |
| Handles the initialization of SAM and ResNet-50 models and manages |
| the end-to-end processing of geographic imagery. |
| """ |
| def __init__(self, config_path="config/config.yaml"): |
| with open(config_path, "r") as f: |
| self.config = yaml.safe_load(f) |
| |
| self.classes = self.config['classes'] |
| self.greenery_classes = self.config['greenery_classes'] |
| |
| |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| print(f"Loading CNN model on {device}...") |
| self.cnn_model = load_model( |
| weights_path=os.path.join(self.config['paths']['output_models'], 'resnet50_eurosat.pth'), |
| num_classes=self.config['model']['num_classes'], |
| device=device |
| ) |
| |
| print("Loading SAM model...") |
| self.sam_generator = get_sam_generator( |
| model_type=self.config['model']['sam_model_type'], |
| checkpoint_path=self.config['paths']['sam_checkpoint'] |
| ) |
| |
| self.estimator = GreeneryEstimator( |
| sam_generator=self.sam_generator, |
| cnn_model=self.cnn_model, |
| greenery_classes=self.greenery_classes, |
| all_classes=self.classes |
| ) |
| print("Pipeline initialized successfully.") |
|
|
| def process_image(self, image_path): |
| """ |
| Runs the end-to-end pipeline on a single image. |
| |
| Args: |
| image_path (str): File path to the input satellite image. |
| |
| Returns: |
| tuple: (image_np, results) where image_np is the RGB image array |
| and results is a dictionary containing greenery metrics and masks. |
| """ |
| image_np = cv2.imread(image_path) |
| if image_np is None: |
| raise FileNotFoundError(f"Could not read image at {image_path} — file may be missing or corrupt") |
| |
| image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB) |
| |
| results = self.estimator.estimate(image_np) |
| return image_np, results |
|
|