ecopulse / src /pipeline.py
acibZ's picture
Deploy EcoPulse
43abac3
Raw
History Blame Contribute Delete
2.34 kB
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']
# Load models
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