geoai-py 0.0.1__py2.py3-none-any.whl → 0.1.5__py2.py3-none-any.whl

geoai/__init__.py

@@ -1,5 +1,5 @@
 """Top-level package for geoai."""
 
 __author__ = """Qiusheng Wu"""
-__email__ = 'giswqs@gmail.com'
-__version__ = '0.0.1'
+__email__ = "giswqs@gmail.com"
+__version__ = "0.1.5"

geoai/common.py

@@ -1,7 +1,6 @@
-"""The common module contains common functions and classes used by the other modules.
-"""
+"""The common module contains common functions and classes used by the other modules."""
+
 
 def hello_world():
-    """Prints "Hello World!" to the console.
-    """
-    print("Hello World!")
+    """Prints "Hello World!" to the console."""
+    print("Hello World!")

geoai/segmentation.py

@@ -0,0 +1,349 @@
+import os
+import numpy as np
+from PIL import Image
+import torch
+import matplotlib.pyplot as plt
+from torch.utils.data import Dataset, Subset
+import torch.nn.functional as F
+from sklearn.model_selection import train_test_split
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    SegformerForSemanticSegmentation,
+    DefaultDataCollator,
+)
+
+
+class CustomDataset(Dataset):
+    """Custom Dataset for loading images and masks."""
+
+    def __init__(
+        self,
+        images_dir: str,
+        masks_dir: str,
+        transform: A.Compose = None,
+        target_size: tuple = (256, 256),
+        num_classes: int = 2,
+    ):
+        """
+        Args:
+            images_dir (str): Directory containing images.
+            masks_dir (str): Directory containing masks.
+            transform (A.Compose, optional): Transformations to be applied on the images and masks.
+            target_size (tuple, optional): Target size for resizing images and masks.
+            num_classes (int, optional): Number of classes in the masks.
+        """
+        self.images_dir = images_dir
+        self.masks_dir = masks_dir
+        self.transform = transform
+        self.target_size = target_size
+        self.num_classes = num_classes
+        self.images = sorted(os.listdir(images_dir))
+        self.masks = sorted(os.listdir(masks_dir))
+
+    def __len__(self) -> int:
+        """Returns the total number of samples."""
+        return len(self.images)
+
+    def __getitem__(self, idx: int) -> dict:
+        """
+        Args:
+            idx (int): Index of the sample to fetch.
+
+        Returns:
+            dict: A dictionary with 'pixel_values' and 'labels'.
+        """
+        img_path = os.path.join(self.images_dir, self.images[idx])
+        mask_path = os.path.join(self.masks_dir, self.masks[idx])
+        image = Image.open(img_path).convert("RGB")
+        mask = Image.open(mask_path).convert("L")
+
+        image = image.resize(self.target_size)
+        mask = mask.resize(self.target_size)
+
+        image = np.array(image)
+        mask = np.array(mask)
+
+        mask = (mask > 127).astype(np.uint8)
+
+        if self.transform:
+            transformed = self.transform(image=image, mask=mask)
+            image = transformed["image"]
+            mask = transformed["mask"]
+
+        assert (
+            mask.max() < self.num_classes
+        ), f"Mask values should be less than {self.num_classes}, but found {mask.max()}"
+        assert (
+            mask.min() >= 0
+        ), f"Mask values should be greater than or equal to 0, but found {mask.min()}"
+
+        mask = mask.clone().detach().long()
+
+        return {"pixel_values": image, "labels": mask}
+
+
+def get_transform() -> A.Compose:
+    """
+    Returns:
+        A.Compose: A composition of image transformations.
+    """
+    return A.Compose(
+        [
+            A.Resize(256, 256),
+            A.HorizontalFlip(p=0.5),
+            A.VerticalFlip(p=0.5),
+            A.RandomRotate90(p=0.5),
+            A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+            ToTensorV2(),
+        ]
+    )
+
+
+def prepare_datasets(
+    images_dir: str,
+    masks_dir: str,
+    transform: A.Compose,
+    test_size: float = 0.2,
+    random_state: int = 42,
+) -> tuple:
+    """
+    Args:
+        images_dir (str): Directory containing images.
+        masks_dir (str): Directory containing masks.
+        transform (A.Compose): Transformations to be applied.
+        test_size (float, optional): Proportion of the dataset to include in the validation split.
+        random_state (int, optional): Random seed for shuffling the dataset.
+
+    Returns:
+        tuple: Training and validation datasets.
+    """
+    dataset = CustomDataset(images_dir, masks_dir, transform)
+    train_indices, val_indices = train_test_split(
+        list(range(len(dataset))), test_size=test_size, random_state=random_state
+    )
+    train_dataset = Subset(dataset, train_indices)
+    val_dataset = Subset(dataset, val_indices)
+    return train_dataset, val_dataset
+
+
+def train_model(
+    train_dataset: Dataset,
+    val_dataset: Dataset,
+    pretrained_model: str = "nvidia/segformer-b0-finetuned-ade-512-512",
+    model_save_path: str = "./model",
+    output_dir: str = "./results",
+    num_epochs: int = 10,
+    batch_size: int = 8,
+    learning_rate: float = 5e-5,
+) -> str:
+    """
+    Trains the model and saves the fine-tuned model to the specified path.
+
+    Args:
+        train_dataset (Dataset): Training dataset.
+        val_dataset (Dataset): Validation dataset.
+        pretrained_model (str, optional): Pretrained model to fine-tune.
+        model_save_path (str): Path to save the fine-tuned model. Defaults to './model'.
+        output_dir (str, optional): Directory to save training outputs.
+        num_epochs (int, optional): Number of training epochs.
+        batch_size (int, optional): Batch size for training and evaluation.
+        learning_rate (float, optional): Learning rate for training.
+
+    Returns:
+        str: Path to the saved fine-tuned model.
+    """
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = SegformerForSemanticSegmentation.from_pretrained(pretrained_model).to(
+        device
+    )
+    data_collator = DefaultDataCollator(return_tensors="pt")
+
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=num_epochs,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=batch_size,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        logging_dir="./logs",
+        learning_rate=learning_rate,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        data_collator=data_collator,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+    )
+
+    trainer.train()
+    model.save_pretrained(model_save_path)
+    print(f"Model saved to {model_save_path}")
+    return model_save_path
+
+
+def load_model(
+    model_path: str, device: torch.device
+) -> SegformerForSemanticSegmentation:
+    """
+    Loads the fine-tuned model from the specified path.
+
+    Args:
+        model_path (str): Path to the model.
+        device (torch.device): Device to load the model on.
+
+    Returns:
+        SegformerForSemanticSegmentation: Loaded model.
+    """
+    model = SegformerForSemanticSegmentation.from_pretrained(model_path)
+    model.to(device)
+    model.eval()
+    return model
+
+
+def preprocess_image(image_path: str, target_size: tuple = (256, 256)) -> torch.Tensor:
+    """
+    Preprocesses the input image for prediction.
+
+    Args:
+        image_path (str): Path to the input image.
+        target_size (tuple, optional): Target size for resizing the image.
+
+    Returns:
+        torch.Tensor: Preprocessed image tensor.
+    """
+    image = Image.open(image_path).convert("RGB")
+    transform = A.Compose(
+        [
+            A.Resize(target_size[0], target_size[1]),
+            A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+            ToTensorV2(),
+        ]
+    )
+    image = np.array(image)
+    transformed = transform(image=image)
+    return transformed["image"].unsqueeze(0)
+
+
+def predict_image(
+    model: SegformerForSemanticSegmentation,
+    image_tensor: torch.Tensor,
+    original_size: tuple,
+    device: torch.device,
+) -> np.ndarray:
+    """
+    Predicts the segmentation mask for the input image.
+
+    Args:
+        model (SegformerForSemanticSegmentation): Fine-tuned model.
+        image_tensor (torch.Tensor): Preprocessed image tensor.
+        original_size (tuple): Original size of the image (width, height).
+        device (torch.device): Device to perform inference on.
+
+    Returns:
+        np.ndarray: Predicted segmentation mask.
+    """
+    with torch.no_grad():
+        image_tensor = image_tensor.to(device)
+        outputs = model(pixel_values=image_tensor)
+        logits = outputs.logits
+        upsampled_logits = F.interpolate(
+            logits, size=original_size[::-1], mode="bilinear", align_corners=False
+        )
+        predictions = torch.argmax(upsampled_logits, dim=1).cpu().numpy()
+    return predictions[0]
+
+
+def segment_image(
+    image_path: str,
+    model_path: str,
+    target_size: tuple = (256, 256),
+    device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+) -> np.ndarray:
+    """
+    Segments the input image using the fine-tuned model.
+
+    Args:
+        image_path (str): Path to the input image.
+        model_path (str): Path to the fine-tuned model.
+        target_size (tuple, optional): Target size for resizing the image.
+        device (torch.device, optional): Device to perform inference on.
+
+    Returns:
+        np.ndarray: Predicted segmentation mask.
+    """
+    model = load_model(model_path, device)
+    image = Image.open(image_path).convert("RGB")
+    original_size = image.size
+    image_tensor = preprocess_image(image_path, target_size)
+    predictions = predict_image(model, image_tensor, original_size, device)
+    return predictions
+
+
+def visualize_predictions(
+    image_path: str,
+    segmented_mask: np.ndarray,
+    target_size: tuple = (256, 256),
+    reference_image_path: str = None,
+) -> None:
+    """
+    Visualizes the original image, segmented mask, and optionally the reference image.
+
+    Args:
+        image_path (str): Path to the original image.
+        segmented_mask (np.ndarray): Predicted segmentation mask.
+        target_size (tuple, optional): Target size for resizing images.
+        reference_image_path (str, optional): Path to the reference image.
+    """
+    original_image = Image.open(image_path).convert("RGB")
+    original_image = original_image.resize(target_size)
+    segmented_image = Image.fromarray((segmented_mask * 255).astype(np.uint8))
+
+    if reference_image_path:
+        reference_image = Image.open(reference_image_path).convert("RGB")
+        reference_image = reference_image.resize(target_size)
+        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+        axes[1].imshow(reference_image)
+        axes[1].set_title("Reference Image")
+        axes[1].axis("off")
+    else:
+        fig, axes = plt.subplots(1, 2, figsize=(12, 6))
+
+    axes[0].imshow(original_image)
+    axes[0].set_title("Original Image")
+    axes[0].axis("off")
+
+    if reference_image_path:
+        axes[2].imshow(segmented_image, cmap="gray")
+        axes[2].set_title("Segmented Image")
+        axes[2].axis("off")
+    else:
+        axes[1].imshow(segmented_image, cmap="gray")
+        axes[1].set_title("Segmented Image")
+        axes[1].axis("off")
+
+    plt.tight_layout()
+    plt.show()
+
+
+# Example usage
+if __name__ == "__main__":
+    images_dir = "../datasets/Water-Bodies-Dataset/Images"
+    masks_dir = "../datasets/Water-Bodies-Dataset/Masks"
+    transform = get_transform()
+    train_dataset, val_dataset = prepare_datasets(images_dir, masks_dir, transform)
+
+    model_save_path = "./fine_tuned_model"
+    train_model(train_dataset, val_dataset, model_save_path)
+
+    image_path = "../datasets/Water-Bodies-Dataset/Images/water_body_44.jpg"
+    reference_image_path = image_path.replace("Images", "Masks")
+    segmented_mask = segment_image(image_path, model_save_path)
+
+    visualize_predictions(
+        image_path, segmented_mask, reference_image_path=reference_image_path
+    )

geoai_py-0.1.5.dist-info/METADATA

@@ -0,0 +1,44 @@
+Metadata-Version: 2.2
+Name: geoai-py
+Version: 0.1.5
+Summary: A Python package for using Artificial Intelligence (AI) with geospatial data
+Author-email: Qiusheng Wu <giswqs@gmail.com>
+License: MIT License
+Project-URL: Homepage, https://github.com/giswqs/geoai
+Keywords: geoai
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: albumentations
+Requires-Dist: scikit-learn
+Requires-Dist: segment-geospatial
+Requires-Dist: torch
+Requires-Dist: transformers
+Provides-Extra: all
+Requires-Dist: geoai[extra]; extra == "all"
+Provides-Extra: extra
+Requires-Dist: pandas; extra == "extra"
+
+# geoai
+
+[![image](https://img.shields.io/pypi/v/geoai-py.svg)](https://pypi.python.org/pypi/geoai-py)
+[![image](https://img.shields.io/conda/vn/conda-forge/geoai.svg)](https://anaconda.org/conda-forge/geoai)
+
+**A Python package for using Artificial Intelligence (AI) with geospatial data**
+
+-   Free software: MIT license
+-   Documentation: <https://geoai.gishub.org>
+
+## Features
+
+-   Visualizing geospatial data, including vector, raster, and LiDAR data
+-   Segmenting remote sensing imagery with the Segment Anything Model
+-   Classifying remote sensing imagery with deep learning models

geoai_py-0.1.5.dist-info/RECORD

@@ -0,0 +1,10 @@
+geoai/__init__.py,sha256=xkyHrnU3iebQB2V0bl1Xd9s8YKI5-UTS-tMbjgnJXNY,120
+geoai/common.py,sha256=Rw6d9qmZDu3dUGTyJto1Y97S7-QA-m2p-pbCNvMDrm4,184
+geoai/geoai.py,sha256=h0hwdogXGFqerm-5ZPeT-irPn91pCcQRjiHThXsRzEk,19
+geoai/segmentation.py,sha256=Vcymnhwl_xikt4v9x8CYJq_vId9R1gB7-YzLfwg-F9M,11372
+geoai_py-0.1.5.dist-info/LICENSE,sha256=vN2L5U7cZ6ZkOHFmc8WiGlsogWsZc5dllMeNxnKVOZg,1070
+geoai_py-0.1.5.dist-info/METADATA,sha256=oQrREZxyg5_OgaVK4Pkn9txkaqUYcUtvigo13MfFNo0,1609
+geoai_py-0.1.5.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109
+geoai_py-0.1.5.dist-info/entry_points.txt,sha256=uGp3Az3HURIsRHP9v-ys0hIbUuBBNUfXv6VbYHIXeg4,41
+geoai_py-0.1.5.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
+geoai_py-0.1.5.dist-info/RECORD,,

{geoai_py-0.0.1.dist-info → geoai_py-0.1.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.1)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py2-none-any
 Tag: py3-none-any

geoai_py-0.1.5.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+geoai = geoai.cli:main

geoai_py-0.0.1.dist-info/METADATA

@@ -1,38 +0,0 @@
-Metadata-Version: 2.1
-Name: geoai-py
-Version: 0.0.1
-Summary: A Python package for using Artificial Intelligence (AI) with geospatial data
-Home-page: https://github.com/opengeos/geoai
-Author: Qiusheng Wu
-Author-email: giswqs@gmail.com
-License: MIT license
-Keywords: geoai
-Classifier: Intended Audience :: Developers
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Natural Language :: English
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-License-File: LICENSE
-
-# geoai
-
-
-[![image](https://img.shields.io/pypi/v/geoai.svg)](https://pypi.python.org/pypi/geoai)
-[![image](https://img.shields.io/conda/vn/conda-forge/geoai.svg)](https://anaconda.org/conda-forge/geoai)
-
-
-**A Python package for using Artificial Intelligence (AI) with geospatial data**
-
-
--   Free software: MIT license
--   Documentation: https://opengeos.github.io/geoai
-    
-
-## Features
-
--   TODO

geoai_py-0.0.1.dist-info/RECORD

@@ -1,8 +0,0 @@
-geoai/__init__.py,sha256=XvWapn8jRuXQLdHwiVTjXdEJErlnJSKLU4hK_2VN3Vc,120
-geoai/common.py,sha256=Si4ZvbkkaVGyKrOa7gmsVNWS2Abtji-V5nEsJhrIM3M,188
-geoai/geoai.py,sha256=h0hwdogXGFqerm-5ZPeT-irPn91pCcQRjiHThXsRzEk,19
-geoai_py-0.0.1.dist-info/LICENSE,sha256=vN2L5U7cZ6ZkOHFmc8WiGlsogWsZc5dllMeNxnKVOZg,1070
-geoai_py-0.0.1.dist-info/METADATA,sha256=9hKtttzpSD-lZc2suU7CX58dWwu6nrtKM6bhHg40AxY,1142
-geoai_py-0.0.1.dist-info/WHEEL,sha256=m9WAupmBd2JGDsXWQGJgMGXIWbQY3F5c2xBJbBhq0nY,110
-geoai_py-0.0.1.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
-geoai_py-0.0.1.dist-info/RECORD,,