File size: 3,824 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 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 | import torch
import numpy as np
from PIL import Image
from src.transforms import EUROSAT_TRANSFORM
class GreeneryEstimator:
"""
Core engine for quantifying vegetation in satellite imagery.
Integrates Instance Segmentation (SAM) and Land-Cover Classification (CNN)
to identify and measure the pixel footprint of greenery.
"""
def __init__(self, sam_generator, cnn_model, greenery_classes, all_classes, device="cuda"):
self.sam_generator = sam_generator
self.cnn_model = cnn_model
self.greenery_classes = greenery_classes
self.all_classes = all_classes
self.device = device if torch.cuda.is_available() else "cpu"
self.cnn_model.to(self.device)
self.cnn_model.eval()
self.transform = EUROSAT_TRANSFORM
def estimate(self, image_np):
"""
Segments an image, classifies individual components, and calculates vegetation density.
Args:
image_np (np.ndarray): The input image in RGB format.
Returns:
dict: A result dictionary containing:
- 'greenery_percentage' (float): Total vegetated area percentage.
- 'green_pixels' (int): Total count of vegetated pixels.
- 'total_pixels' (int): Total image resolution.
- 'mask_classifications' (list): Detailed metadata for every segmented object.
"""
# 1. Generate SAM masks
masks = self.sam_generator.generate(image_np)
total_pixels = image_np.shape[0] * image_np.shape[1]
green_pixels = 0
mask_classifications = []
# 2. Iterate and classify each mask
for idx, mask_data in enumerate(masks):
bbox = mask_data['bbox'] # [x, y, w, h]
x, y, w, h = [int(v) for v in bbox]
h_img, w_img = image_np.shape[:2]
# Clamp bounding box coordinates to image boundaries
# to prevent silently wrong crop sizes from out-of-bounds masks
x = max(0, min(x, w_img - 1))
y = max(0, min(y, h_img - 1))
w = min(w, w_img - x)
h = min(h, h_img - y)
# Avoid completely empty or tiny boxes
if w <= 0 or h <= 0:
continue
# Crop the original image using the bounding box
crop_np = image_np[y:y+h, x:x+w]
if crop_np.size == 0:
continue
crop_img = Image.fromarray(crop_np)
input_tensor = self.transform(crop_img).unsqueeze(0).to(self.device)
with torch.no_grad():
output = self.cnn_model(input_tensor)
_, predicted = torch.max(output, 1)
predicted_class = self.all_classes[predicted.item()]
# Check if classified as greenery
is_green = predicted_class in self.greenery_classes
# Count pixels strictly within the boolean mask, not just the bounding box
segmentation = mask_data['segmentation']
mask_pixels = np.sum(segmentation)
if is_green:
green_pixels += mask_pixels
mask_classifications.append({
'mask_id': idx,
'class': predicted_class,
'is_green': is_green,
'pixels': mask_pixels,
'segmentation': segmentation,
'bbox': bbox
})
greenery_percentage = (green_pixels / total_pixels) * 100
return {
'greenery_percentage': greenery_percentage,
'green_pixels': green_pixels,
'total_pixels': total_pixels,
'mask_classifications': mask_classifications
}
|