File size: 2,337 Bytes
43abac3 | 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 | 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
|