geoai-py 0.8.1__py2.py3-none-any.whl → 0.8.3__py2.py3-none-any.whl

geoai/__init__.py

@@ -2,7 +2,7 @@
 
 __author__ = """Qiusheng Wu"""
 __email__ = "giswqs@gmail.com"
-__version__ = "0.8.1"
+__version__ = "0.8.3"
 
 
 import os

geoai/geoai.py

@@ -32,6 +32,11 @@ from .train import (
     train_segmentation_model,
     semantic_segmentation,
     semantic_segmentation_batch,
+    train_instance_segmentation_model,
+    instance_segmentation,
+    instance_segmentation_batch,
+    get_instance_segmentation_model,
+    instance_segmentation_inference_on_geotiff,
 )
 from .utils import *
 

geoai/train.py

@@ -750,6 +750,15 @@ def train_MaskRCNN_model(
     start_epoch = 0
     best_iou = 0
 
+    # Initialize training history
+    training_history = {
+        "train_loss": [],
+        "val_loss": [],
+        "val_iou": [],
+        "epochs": [],
+        "lr": [],
+    }
+
     # Load pretrained model if provided
     if pretrained_model_path:
         if not os.path.exists(pretrained_model_path):
@@ -800,6 +809,13 @@ def train_MaskRCNN_model(
         # Evaluate
         eval_metrics = evaluate(model, val_loader, device)
 
+        # Record training history
+        training_history["train_loss"].append(train_loss)
+        training_history["val_loss"].append(eval_metrics["loss"])
+        training_history["val_iou"].append(eval_metrics["IoU"])
+        training_history["epochs"].append(epoch + 1)
+        training_history["lr"].append(optimizer.param_groups[0]["lr"])
+
         # Print metrics
         print(
             f"Epoch {epoch+1}/{num_epochs}: Train Loss: {train_loss:.4f}, Val Loss: {eval_metrics['loss']:.4f}, Val IoU: {eval_metrics['IoU']:.4f}"
@@ -811,33 +827,11 @@ def train_MaskRCNN_model(
             print(f"Saving best model with IoU: {best_iou:.4f}")
             torch.save(model.state_dict(), os.path.join(output_dir, "best_model.pth"))
 
-        # Save checkpoint every 10 epochs
-        if (epoch + 1) % 10 == 0 or epoch == num_epochs - 1:
-            torch.save(
-                {
-                    "epoch": epoch,
-                    "model_state_dict": model.state_dict(),
-                    "optimizer_state_dict": optimizer.state_dict(),
-                    "scheduler_state_dict": lr_scheduler.state_dict(),
-                    "best_iou": best_iou,
-                },
-                os.path.join(output_dir, f"checkpoint_epoch_{epoch+1}.pth"),
-            )
-
     # Save final model
     torch.save(model.state_dict(), os.path.join(output_dir, "final_model.pth"))
 
-    # Save full checkpoint of final state
-    torch.save(
-        {
-            "epoch": num_epochs - 1,
-            "model_state_dict": model.state_dict(),
-            "optimizer_state_dict": optimizer.state_dict(),
-            "scheduler_state_dict": lr_scheduler.state_dict(),
-            "best_iou": best_iou,
-        },
-        os.path.join(output_dir, "final_checkpoint.pth"),
-    )
+    # Save training history
+    torch.save(training_history, os.path.join(output_dir, "training_history.pth"))
 
     # Load best model for evaluation and visualization
     model.load_state_dict(torch.load(os.path.join(output_dir, "best_model.pth")))
@@ -1101,6 +1095,237 @@ def inference_on_geotiff(
         return output_path, inference_time
 
 
+def instance_segmentation_inference_on_geotiff(
+    model,
+    geotiff_path,
+    output_path,
+    window_size=512,
+    overlap=256,
+    confidence_threshold=0.5,
+    batch_size=4,
+    num_channels=3,
+    device=None,
+    **kwargs,
+):
+    """
+    Perform instance segmentation inference on a large GeoTIFF using a sliding window approach.
+
+    This function collects all detections first, then applies non-maximum suppression
+    to handle overlapping detections from different windows, preventing artifacts.
+
+    Args:
+        model (torch.nn.Module): Trained model for inference.
+        geotiff_path (str): Path to input GeoTIFF file.
+        output_path (str): Path to save output instance mask GeoTIFF.
+        window_size (int): Size of sliding window for inference.
+        overlap (int): Overlap between adjacent windows.
+        confidence_threshold (float): Confidence threshold for predictions (0-1).
+        batch_size (int): Batch size for inference.
+        num_channels (int): Number of channels to use from the input image.
+        device (torch.device, optional): Device to run inference on. If None, uses CUDA if available.
+        **kwargs: Additional arguments.
+
+    Returns:
+        tuple: Tuple containing output path and inference time in seconds.
+    """
+    if device is None:
+        device = get_device()
+
+    # Put model in evaluation mode
+    model.to(device)
+    model.eval()
+
+    # Open the GeoTIFF
+    with rasterio.open(geotiff_path) as src:
+        # Read metadata
+        meta = src.meta
+        height = src.height
+        width = src.width
+
+        # Update metadata for output raster
+        out_meta = meta.copy()
+        out_meta.update(
+            {"count": 1, "dtype": "uint16"}  # uint16 to support many instances
+        )
+
+        # Store all detections globally for NMS
+        all_detections = []
+
+        # Calculate the number of windows needed to cover the entire image
+        steps_y = math.ceil((height - overlap) / (window_size - overlap))
+        steps_x = math.ceil((width - overlap) / (window_size - overlap))
+
+        # Ensure we cover the entire image
+        last_y = height - window_size
+        last_x = width - window_size
+
+        total_windows = steps_y * steps_x
+        print(
+            f"Processing {total_windows} windows with size {window_size}x{window_size} and overlap {overlap}..."
+        )
+
+        # Create progress bar
+        pbar = tqdm(total=total_windows)
+
+        # Process in batches
+        batch_inputs = []
+        batch_positions = []
+        batch_count = 0
+
+        start_time = time.time()
+
+        # Slide window over the image
+        for i in range(steps_y + 1):  # +1 to ensure we reach the edge
+            y = min(i * (window_size - overlap), last_y)
+            y = max(0, y)  # Prevent negative indices
+
+            if y > last_y and i > 0:  # Skip if we've already covered the entire height
+                continue
+
+            for j in range(steps_x + 1):  # +1 to ensure we reach the edge
+                x = min(j * (window_size - overlap), last_x)
+                x = max(0, x)  # Prevent negative indices
+
+                if (
+                    x > last_x and j > 0
+                ):  # Skip if we've already covered the entire width
+                    continue
+
+                # Read window
+                window = src.read(window=Window(x, y, window_size, window_size))
+
+                # Check if window is valid
+                if window.shape[1] == 0 or window.shape[2] == 0:
+                    continue
+
+                # Handle edge cases where window might be smaller than expected
+                actual_height, actual_width = window.shape[1], window.shape[2]
+
+                # Convert to [C, H, W] format and normalize
+                image = window.astype(np.float32) / 255.0
+
+                # Handle different number of channels
+                if image.shape[0] > num_channels:
+                    image = image[:num_channels]
+                elif image.shape[0] < num_channels:
+                    # Pad with zeros if less than expected channels
+                    padded = np.zeros(
+                        (num_channels, image.shape[1], image.shape[2]), dtype=np.float32
+                    )
+                    padded[: image.shape[0]] = image
+                    image = padded
+
+                # Convert to tensor
+                image_tensor = torch.tensor(image, device=device)
+
+                # Add to batch
+                batch_inputs.append(image_tensor)
+                batch_positions.append((y, x, actual_height, actual_width))
+                batch_count += 1
+
+                # Process batch when it reaches the batch size or at the end
+                if batch_count == batch_size or (i == steps_y and j == steps_x):
+                    # Forward pass
+                    with torch.no_grad():
+                        outputs = model(batch_inputs)
+
+                    # Process each output in the batch
+                    for idx, output in enumerate(outputs):
+                        y_pos, x_pos, h, w = batch_positions[idx]
+
+                        # Process each detected instance
+                        if len(output["scores"]) > 0:
+                            # Get instances that meet confidence threshold
+                            keep = output["scores"] > confidence_threshold
+                            masks = output["masks"][keep].squeeze(1)
+                            scores = output["scores"][keep]
+                            boxes = output["boxes"][keep]
+
+                            # Convert to global coordinates and store
+                            for k in range(len(masks)):
+                                mask = masks[k].cpu().numpy() > 0.5
+                                score = scores[k].cpu().item()
+                                box = boxes[k].cpu().numpy()
+
+                                # Convert box to global coordinates
+                                global_box = [
+                                    box[0] + x_pos,
+                                    box[1] + y_pos,
+                                    box[2] + x_pos,
+                                    box[3] + y_pos,
+                                ]
+
+                                # Create global mask
+                                global_mask = np.zeros((height, width), dtype=bool)
+                                global_mask[y_pos : y_pos + h, x_pos : x_pos + w] = mask
+
+                                all_detections.append(
+                                    {
+                                        "mask": global_mask,
+                                        "score": score,
+                                        "box": global_box,
+                                    }
+                                )
+
+                    # Reset batch
+                    batch_inputs = []
+                    batch_positions = []
+                    batch_count = 0
+
+                    # Update progress bar
+                    pbar.update(len(outputs))
+
+        # Close progress bar
+        pbar.close()
+
+        print(f"Collected {len(all_detections)} detections before NMS")
+
+        # Apply Non-Maximum Suppression to handle overlapping detections
+        if len(all_detections) > 0:
+            # Convert to tensors for NMS
+            boxes = torch.tensor([det["box"] for det in all_detections])
+            scores = torch.tensor([det["score"] for det in all_detections])
+
+            # Apply NMS with IoU threshold
+            nms_threshold = 0.3  # IoU threshold for NMS
+            keep_indices = torchvision.ops.nms(boxes, scores, nms_threshold)
+
+            # Keep only the selected detections
+            final_detections = [all_detections[i] for i in keep_indices]
+            print(f"After NMS: {len(final_detections)} detections")
+
+            # Create final instance mask
+            instance_mask = np.zeros((height, width), dtype=np.uint16)
+
+            # Sort by score (highest first) for consistent ordering
+            final_detections.sort(key=lambda x: x["score"], reverse=True)
+
+            # Assign unique IDs to each detection
+            for instance_id, detection in enumerate(final_detections, 1):
+                mask = detection["mask"]
+                # Only assign to pixels that are not already assigned
+                available_pixels = (instance_mask == 0) & mask
+                instance_mask[available_pixels] = instance_id
+        else:
+            # No detections found
+            instance_mask = np.zeros((height, width), dtype=np.uint16)
+
+        # Record time
+        inference_time = time.time() - start_time
+        print(f"Instance segmentation completed in {inference_time:.2f} seconds")
+        print(
+            f"Final instances: {len(final_detections) if len(all_detections) > 0 else 0}"
+        )
+
+        # Save output
+        with rasterio.open(output_path, "w", **out_meta) as dst:
+            dst.write(instance_mask, 1)
+
+        print(f"Saved instance segmentation to {output_path}")
+
+        return output_path, inference_time
+
+
 def object_detection(
     input_path,
     output_path,
@@ -3067,3 +3292,220 @@ def semantic_segmentation_batch(
             continue
 
     print(f"Batch processing completed. Results saved to {output_dir}")
+
+
+def train_instance_segmentation_model(
+    images_dir,
+    labels_dir,
+    output_dir,
+    num_classes=2,
+    num_channels=3,
+    batch_size=4,
+    num_epochs=10,
+    learning_rate=0.005,
+    seed=42,
+    val_split=0.2,
+    visualize=False,
+    device=None,
+    verbose=True,
+    **kwargs,
+):
+    """
+    Train an instance segmentation model using Mask R-CNN.
+
+    This is a wrapper function for train_MaskRCNN_model with clearer naming.
+
+    Args:
+        images_dir (str): Directory containing image GeoTIFF files.
+        labels_dir (str): Directory containing label GeoTIFF files.
+        output_dir (str): Directory to save model checkpoints and results.
+        num_classes (int): Number of classes (including background). Defaults to 2.
+        num_channels (int): Number of input channels. Defaults to 3.
+        batch_size (int): Batch size for training. Defaults to 4.
+        num_epochs (int): Number of training epochs. Defaults to 10.
+        learning_rate (float): Initial learning rate. Defaults to 0.005.
+        seed (int): Random seed for reproducibility. Defaults to 42.
+        val_split (float): Fraction of data to use for validation (0-1). Defaults to 0.2.
+        visualize (bool): Whether to generate visualizations. Defaults to False.
+        device (torch.device): Device to train on. If None, uses CUDA if available.
+        verbose (bool): If True, prints detailed training progress. Defaults to True.
+        **kwargs: Additional arguments passed to train_MaskRCNN_model.
+
+    Returns:
+        None: Model weights are saved to output_dir.
+    """
+    # Create model with the specified number of classes
+    model = get_instance_segmentation_model(
+        num_classes=num_classes, num_channels=num_channels, pretrained=True
+    )
+
+    return train_MaskRCNN_model(
+        images_dir=images_dir,
+        labels_dir=labels_dir,
+        output_dir=output_dir,
+        num_channels=num_channels,
+        model=model,
+        batch_size=batch_size,
+        num_epochs=num_epochs,
+        learning_rate=learning_rate,
+        seed=seed,
+        val_split=val_split,
+        visualize=visualize,
+        device=device,
+        verbose=verbose,
+        **kwargs,
+    )
+
+
+def instance_segmentation(
+    input_path,
+    output_path,
+    model_path,
+    window_size=512,
+    overlap=256,
+    confidence_threshold=0.5,
+    batch_size=4,
+    num_channels=3,
+    num_classes=2,
+    device=None,
+    **kwargs,
+):
+    """
+    Perform instance segmentation on a GeoTIFF using a pre-trained Mask R-CNN model.
+
+    This is a wrapper function for object_detection with clearer naming.
+
+    Args:
+        input_path (str): Path to input GeoTIFF file.
+        output_path (str): Path to save output mask GeoTIFF.
+        model_path (str): Path to trained model weights.
+        window_size (int): Size of sliding window for inference. Defaults to 512.
+        overlap (int): Overlap between adjacent windows. Defaults to 256.
+        confidence_threshold (float): Confidence threshold for predictions (0-1). Defaults to 0.5.
+        batch_size (int): Batch size for inference. Defaults to 4.
+        num_channels (int): Number of channels in the input image and model. Defaults to 3.
+        num_classes (int): Number of classes (including background). Defaults to 2.
+        device (torch.device): Device to run inference on. If None, uses CUDA if available.
+        **kwargs: Additional arguments passed to object_detection.
+
+    Returns:
+        None: Output mask is saved to output_path.
+    """
+    # Create model with the specified number of classes
+    model = get_instance_segmentation_model(
+        num_classes=num_classes, num_channels=num_channels, pretrained=True
+    )
+
+    # Load the trained model
+    if device is None:
+        device = get_device()
+
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.to(device)
+
+    # Use the proper instance segmentation inference function
+    return instance_segmentation_inference_on_geotiff(
+        model=model,
+        geotiff_path=input_path,
+        output_path=output_path,
+        window_size=window_size,
+        overlap=overlap,
+        confidence_threshold=confidence_threshold,
+        batch_size=batch_size,
+        num_channels=num_channels,
+        device=device,
+        **kwargs,
+    )
+
+
+def instance_segmentation_batch(
+    input_dir,
+    output_dir,
+    model_path,
+    window_size=512,
+    overlap=256,
+    confidence_threshold=0.5,
+    batch_size=4,
+    num_channels=3,
+    num_classes=2,
+    device=None,
+    **kwargs,
+):
+    """
+    Perform instance segmentation on multiple GeoTIFF files using a pre-trained Mask R-CNN model.
+
+    This is a wrapper function for object_detection_batch with clearer naming.
+
+    Args:
+        input_dir (str): Directory containing input GeoTIFF files.
+        output_dir (str): Directory to save output mask GeoTIFF files.
+        model_path (str): Path to trained model weights.
+        window_size (int): Size of sliding window for inference. Defaults to 512.
+        overlap (int): Overlap between adjacent windows. Defaults to 256.
+        confidence_threshold (float): Confidence threshold for predictions (0-1). Defaults to 0.5.
+        batch_size (int): Batch size for inference. Defaults to 4.
+        num_channels (int): Number of channels in the input image and model. Defaults to 3.
+        num_classes (int): Number of classes (including background). Defaults to 2.
+        device (torch.device): Device to run inference on. If None, uses CUDA if available.
+        **kwargs: Additional arguments passed to object_detection_batch.
+
+    Returns:
+        None: Output masks are saved to output_dir.
+    """
+    # Create model with the specified number of classes
+    model = get_instance_segmentation_model(
+        num_classes=num_classes, num_channels=num_channels, pretrained=True
+    )
+
+    # Load the trained model
+    if device is None:
+        device = get_device()
+
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.to(device)
+
+    # Process all GeoTIFF files in the input directory
+    import glob
+
+    input_files = glob.glob(os.path.join(input_dir, "*.tif")) + glob.glob(
+        os.path.join(input_dir, "*.tiff")
+    )
+
+    if not input_files:
+        print(f"No GeoTIFF files found in {input_dir}")
+        return
+
+    # Create output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(f"Processing {len(input_files)} files...")
+
+    for input_file in input_files:
+        try:
+            # Generate output filename
+            base_name = os.path.splitext(os.path.basename(input_file))[0]
+            output_file = os.path.join(output_dir, f"{base_name}_instances.tif")
+
+            print(f"Processing {input_file}...")
+
+            # Run instance segmentation inference
+            instance_segmentation_inference_on_geotiff(
+                model=model,
+                geotiff_path=input_file,
+                output_path=output_file,
+                window_size=window_size,
+                overlap=overlap,
+                confidence_threshold=confidence_threshold,
+                batch_size=batch_size,
+                num_channels=num_channels,
+                device=device,
+                **kwargs,
+            )
+
+            print(f"Saved result to {output_file}")
+
+        except Exception as e:
+            print(f"Error processing {input_file}: {str(e)}")
+            continue
+
+    print(f"Batch processing completed. Results saved to {output_dir}")

geoai/utils.py

@@ -157,7 +157,8 @@ def view_image(
     image: Union[np.ndarray, torch.Tensor],
     transpose: bool = False,
     bdx: Optional[int] = None,
-    scale_factor: float = 1.0,
+    clip_percentiles: Optional[Tuple[float, float]] = (2, 98),
+    gamma: Optional[float] = None,
     figsize: Tuple[int, int] = (10, 5),
     axis_off: bool = True,
     title: Optional[str] = None,
@@ -185,7 +186,7 @@ def view_image(
     elif isinstance(image, str):
         image = rasterio.open(image).read().transpose(1, 2, 0)
 
-    plt.figure(figsize=figsize)
+    ax = plt.figure(figsize=figsize)
 
     if transpose:
         image = image.transpose(1, 2, 0)
@@ -196,8 +197,14 @@ def view_image(
     if len(image.shape) > 2 and image.shape[2] > 3:
         image = image[:, :, 0:3]
 
-    if scale_factor != 1.0:
-        image = np.clip(image * scale_factor, 0, 1)
+    if clip_percentiles is not None:
+        p_low, p_high = clip_percentiles
+        lower = np.percentile(image, p_low)
+        upper = np.percentile(image, p_high)
+        image = np.clip((image - lower) / (upper - lower), 0, 1)
+
+    if gamma is not None:
+        image = np.power(image, gamma)
 
     plt.imshow(image, **kwargs)
     if axis_off:
@@ -207,6 +214,8 @@ def view_image(
     plt.show()
     plt.close()
 
+    return ax
+
 
 def plot_images(
     images: Iterable[torch.Tensor],
@@ -5658,7 +5667,7 @@ def orthogonalize(
         if len(df) == 0:
             return ring
 
-        # If we have a triangle-like result (3 segments), return the original shape
+        # If we have a triangle-like result (3 segments or less), return the original shape
         if len(df) <= 3:
             return ring
 
@@ -5669,8 +5678,116 @@ def orthogonalize(
         if len(joined_ring) == 0 or len(joined_ring[0]) < 3:
             return ring
 
-        # Basic validation: if result has 3 or fewer points (triangle), use original
-        if len(joined_ring[0]) <= 3:
+        # Enhanced validation: check for triangular result and geometric validity
+        result_coords = joined_ring[0]
+
+        # If result has 3 or fewer points (triangle), use original
+        if len(result_coords) <= 3:  # 2 points + closing point (degenerate)
+            return ring
+
+        # Additional validation: check for degenerate geometry
+        # Calculate area ratio to detect if the shape got severely distorted
+        def calculate_polygon_area(coords):
+            if len(coords) < 3:
+                return 0
+            area = 0
+            n = len(coords)
+            for i in range(n):
+                j = (i + 1) % n
+                area += coords[i][0] * coords[j][1]
+                area -= coords[j][0] * coords[i][1]
+            return abs(area) / 2
+
+        original_area = calculate_polygon_area(ring)
+        result_area = calculate_polygon_area(result_coords)
+
+        # If the area changed dramatically (more than 30% shrinkage or 300% growth), use original
+        if original_area > 0 and result_area > 0:
+            area_ratio = result_area / original_area
+            if area_ratio < 0.3 or area_ratio > 3.0:
+                return ring
+
+        # Check for triangular spikes and problematic artifacts
+        very_acute_angle_count = 0
+        triangular_spike_detected = False
+
+        for i in range(len(result_coords) - 1):  # -1 to exclude closing point
+            p1 = result_coords[i - 1]
+            p2 = result_coords[i]
+            p3 = result_coords[(i + 1) % (len(result_coords) - 1)]
+
+            # Calculate angle at p2
+            v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
+            v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
+
+            v1_norm = np.linalg.norm(v1)
+            v2_norm = np.linalg.norm(v2)
+
+            if v1_norm > 0 and v2_norm > 0:
+                cos_angle = np.dot(v1, v2) / (v1_norm * v2_norm)
+                cos_angle = np.clip(cos_angle, -1, 1)
+                angle = np.arccos(cos_angle)
+
+                # Count very acute angles (< 20 degrees) - these are likely spikes
+                if angle < np.pi / 9:  # 20 degrees
+                    very_acute_angle_count += 1
+                    # If it's very acute with short sides, it's definitely a spike
+                    if v1_norm < 5 or v2_norm < 5:
+                        triangular_spike_detected = True
+
+        # Check for excessively long edges that might be artifacts
+        edge_lengths = []
+        for i in range(len(result_coords) - 1):
+            edge_len = np.sqrt(
+                (result_coords[i + 1][0] - result_coords[i][0]) ** 2
+                + (result_coords[i + 1][1] - result_coords[i][1]) ** 2
+            )
+            edge_lengths.append(edge_len)
+
+        excessive_edge_detected = False
+        if len(edge_lengths) > 0:
+            avg_edge_length = np.mean(edge_lengths)
+            max_edge_length = np.max(edge_lengths)
+            # Only reject if edge is extremely disproportionate (8x average)
+            if max_edge_length > avg_edge_length * 8:
+                excessive_edge_detected = True
+
+        # Check for triangular artifacts by detecting spikes that extend beyond bounds
+        # Calculate original bounds
+        orig_xs = [p[0] for p in ring]
+        orig_ys = [p[1] for p in ring]
+        orig_min_x, orig_max_x = min(orig_xs), max(orig_xs)
+        orig_min_y, orig_max_y = min(orig_ys), max(orig_ys)
+        orig_width = orig_max_x - orig_min_x
+        orig_height = orig_max_y - orig_min_y
+
+        # Calculate result bounds
+        result_xs = [p[0] for p in result_coords]
+        result_ys = [p[1] for p in result_coords]
+        result_min_x, result_max_x = min(result_xs), max(result_xs)
+        result_min_y, result_max_y = min(result_ys), max(result_ys)
+
+        # Stricter bounds checking to catch triangular artifacts
+        bounds_extension_detected = False
+        # More conservative: only allow 10% extension
+        tolerance_x = max(orig_width * 0.1, 1.0)  # 10% tolerance, at least 1 unit
+        tolerance_y = max(orig_height * 0.1, 1.0)  # 10% tolerance, at least 1 unit
+
+        if (
+            result_min_x < orig_min_x - tolerance_x
+            or result_max_x > orig_max_x + tolerance_x
+            or result_min_y < orig_min_y - tolerance_y
+            or result_max_y > orig_max_y + tolerance_y
+        ):
+            bounds_extension_detected = True
+
+        # Reject if we detect triangular spikes, excessive edges, or bounds violations
+        if (
+            triangular_spike_detected
+            or very_acute_angle_count > 2  # Multiple very acute angles
+            or excessive_edge_detected
+            or bounds_extension_detected
+        ):  # Any significant bounds extension
             return ring
 
         # Convert back to a list and ensure it's closed
@@ -5887,37 +6004,86 @@ def orthogonalize(
                 }
             )
 
+        # Improved fix: Prevent merging that would create triangular or problematic shapes
         if (
-            len(ortho_list) > 0 and ortho_list[0]["angle"] == ortho_list[-1]["angle"]
+            len(ortho_list) > 3 and ortho_list[0]["angle"] == ortho_list[-1]["angle"]
         ):  # join first and last segment if they're in same direction
-            totlen = ortho_list[0]["len"] + ortho_list[-1]["len"]
-            merge_cx = (
-                (ortho_list[0]["cx"] * ortho_list[0]["len"])
-                + (ortho_list[-1]["cx"] * ortho_list[-1]["len"])
-            ) / totlen
-
-            merge_cy = (
-                (ortho_list[0]["cy"] * ortho_list[0]["len"])
-                + (ortho_list[-1]["cy"] * ortho_list[-1]["len"])
-            ) / totlen
-
-            rot_angle = ortho_list[0]["angle"]
-            X1 = merge_cx - (totlen / 2) * math.cos(rot_angle)
-            X2 = merge_cx + (totlen / 2) * math.cos(rot_angle)
-            Y1 = merge_cy - (totlen / 2) * math.sin(rot_angle)
-            Y2 = merge_cy + (totlen / 2) * math.sin(rot_angle)
-
-            ortho_list[-1] = {
-                "x1": X1,
-                "y1": Y1,
-                "x2": X2,
-                "y2": Y2,
-                "len": totlen,
-                "cx": merge_cx,
-                "cy": merge_cy,
-                "angle": rot_angle,
-            }
-            ortho_list = ortho_list[1:]
+            # Check if merging would result in 3 or 4 segments (potentially triangular)
+            resulting_segments = len(ortho_list) - 1
+            if resulting_segments <= 4:
+                # For very small polygons, be extra cautious about merging
+                # Calculate the spatial relationship between first and last segments
+                first_center = np.array([ortho_list[0]["cx"], ortho_list[0]["cy"]])
+                last_center = np.array([ortho_list[-1]["cx"], ortho_list[-1]["cy"]])
+                center_distance = np.linalg.norm(first_center - last_center)
+
+                # Get average segment length for comparison
+                avg_length = sum(seg["len"] for seg in ortho_list) / len(ortho_list)
+
+                # Only merge if segments are close enough and it won't create degenerate shapes
+                if center_distance > avg_length * 1.5:
+                    # Skip merging - segments are too far apart
+                    pass
+                else:
+                    # Proceed with merging only for well-connected segments
+                    totlen = ortho_list[0]["len"] + ortho_list[-1]["len"]
+                    merge_cx = (
+                        (ortho_list[0]["cx"] * ortho_list[0]["len"])
+                        + (ortho_list[-1]["cx"] * ortho_list[-1]["len"])
+                    ) / totlen
+
+                    merge_cy = (
+                        (ortho_list[0]["cy"] * ortho_list[0]["len"])
+                        + (ortho_list[-1]["cy"] * ortho_list[-1]["len"])
+                    ) / totlen
+
+                    rot_angle = ortho_list[0]["angle"]
+                    X1 = merge_cx - (totlen / 2) * math.cos(rot_angle)
+                    X2 = merge_cx + (totlen / 2) * math.cos(rot_angle)
+                    Y1 = merge_cy - (totlen / 2) * math.sin(rot_angle)
+                    Y2 = merge_cy + (totlen / 2) * math.sin(rot_angle)
+
+                    ortho_list[-1] = {
+                        "x1": X1,
+                        "y1": Y1,
+                        "x2": X2,
+                        "y2": Y2,
+                        "len": totlen,
+                        "cx": merge_cx,
+                        "cy": merge_cy,
+                        "angle": rot_angle,
+                    }
+                    ortho_list = ortho_list[1:]
+            else:
+                # For larger polygons, proceed with standard merging
+                totlen = ortho_list[0]["len"] + ortho_list[-1]["len"]
+                merge_cx = (
+                    (ortho_list[0]["cx"] * ortho_list[0]["len"])
+                    + (ortho_list[-1]["cx"] * ortho_list[-1]["len"])
+                ) / totlen
+
+                merge_cy = (
+                    (ortho_list[0]["cy"] * ortho_list[0]["len"])
+                    + (ortho_list[-1]["cy"] * ortho_list[-1]["len"])
+                ) / totlen
+
+                rot_angle = ortho_list[0]["angle"]
+                X1 = merge_cx - (totlen / 2) * math.cos(rot_angle)
+                X2 = merge_cx + (totlen / 2) * math.cos(rot_angle)
+                Y1 = merge_cy - (totlen / 2) * math.sin(rot_angle)
+                Y2 = merge_cy + (totlen / 2) * math.sin(rot_angle)
+
+                ortho_list[-1] = {
+                    "x1": X1,
+                    "y1": Y1,
+                    "x2": X2,
+                    "y2": Y2,
+                    "len": totlen,
+                    "cx": merge_cx,
+                    "cy": merge_cy,
+                    "angle": rot_angle,
+                }
+                ortho_list = ortho_list[1:]
         ortho_df = pd.DataFrame(ortho_list)
         return ortho_df
 
@@ -6026,12 +6192,49 @@ def orthogonalize(
                     np.sqrt((x4 - x3) ** 2 + (y4 - y3) ** 2),
                 )
 
-                # If intersection is too far away, use the endpoint of the first segment instead
-                if dist_to_seg1 > max_len * 0.5 or dist_to_seg2 > max_len * 0.5:
-                    ring.append([x2, y2])
+                # Improved intersection validation
+                # Calculate angle between segments to detect sharp corners
+                v1 = np.array([x2 - x1, y2 - y1])
+                v2 = np.array([x4 - x3, y4 - y3])
+                v1_norm = np.linalg.norm(v1)
+                v2_norm = np.linalg.norm(v2)
+
+                if v1_norm > 0 and v2_norm > 0:
+                    cos_angle = np.dot(v1, v2) / (v1_norm * v2_norm)
+                    cos_angle = np.clip(cos_angle, -1, 1)
+                    angle = np.arccos(cos_angle)
+
+                    # Check for very sharp angles that could create triangular artifacts
+                    is_sharp_angle = (
+                        angle < np.pi / 6 or angle > 5 * np.pi / 6
+                    )  # <30° or >150°
+                else:
+                    is_sharp_angle = False
+
+                # Determine whether to use intersection or segment endpoint
+                if (
+                    dist_to_seg1 > max_len * 0.5
+                    or dist_to_seg2 > max_len * 0.5
+                    or is_sharp_angle
+                ):
+                    # Use a more conservative approach for problematic intersections
+                    # Use the closer endpoint between segments
+                    dist_x2_to_seg2 = min(
+                        np.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2),
+                        np.sqrt((x2 - x4) ** 2 + (y2 - y4) ** 2),
+                    )
+                    dist_x3_to_seg1 = min(
+                        np.sqrt((x3 - x1) ** 2 + (y3 - y1) ** 2),
+                        np.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2),
+                    )
+
+                    if dist_x2_to_seg2 <= dist_x3_to_seg1:
+                        ring.append([x2, y2])
+                    else:
+                        ring.append([x3, y3])
                 else:
                     ring.append(intersection)
-            except Exception as e:
+            except Exception:
                 # If intersection calculation fails, use the endpoint of the first segment
                 ring.append([x2, y2])
 
@@ -6057,11 +6260,42 @@ def orthogonalize(
                 np.sqrt((x4 - x3) ** 2 + (y4 - y3) ** 2),
             )
 
-            if dist_to_seg1 > max_len * 0.5 or dist_to_seg2 > max_len * 0.5:
-                ring.append([x2, y2])
+            # Apply same sharp angle detection for closing segment
+            v1 = np.array([x2 - x1, y2 - y1])
+            v2 = np.array([x4 - x3, y4 - y3])
+            v1_norm = np.linalg.norm(v1)
+            v2_norm = np.linalg.norm(v2)
+
+            if v1_norm > 0 and v2_norm > 0:
+                cos_angle = np.dot(v1, v2) / (v1_norm * v2_norm)
+                cos_angle = np.clip(cos_angle, -1, 1)
+                angle = np.arccos(cos_angle)
+                is_sharp_angle = angle < np.pi / 6 or angle > 5 * np.pi / 6
+            else:
+                is_sharp_angle = False
+
+            if (
+                dist_to_seg1 > max_len * 0.5
+                or dist_to_seg2 > max_len * 0.5
+                or is_sharp_angle
+            ):
+                # Use conservative approach for closing segment
+                dist_x2_to_seg2 = min(
+                    np.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2),
+                    np.sqrt((x2 - x4) ** 2 + (y2 - y4) ** 2),
+                )
+                dist_x3_to_seg1 = min(
+                    np.sqrt((x3 - x1) ** 2 + (y3 - y1) ** 2),
+                    np.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2),
+                )
+
+                if dist_x2_to_seg2 <= dist_x3_to_seg1:
+                    ring.append([x2, y2])
+                else:
+                    ring.append([x3, y3])
             else:
                 ring.append(intersection)
-        except Exception as e:
+        except Exception:
             # If intersection calculation fails, use the endpoint of the last segment
             ring.append([x2, y2])
 
@@ -7219,41 +7453,61 @@ def plot_performance_metrics(history_path, figsize=(15, 5), verbose=True):
     """
     history = torch.load(history_path)
 
+    # Handle different key naming conventions
+    train_loss_key = "train_losses" if "train_losses" in history else "train_loss"
+    val_loss_key = "val_losses" if "val_losses" in history else "val_loss"
+    val_iou_key = "val_ious" if "val_ious" in history else "val_iou"
+    val_dice_key = "val_dices" if "val_dices" in history else "val_dice"
+
+    # Determine number of subplots based on available metrics
+    has_dice = val_dice_key in history
+    n_plots = 3 if has_dice else 2
+    figsize = (15, 5) if has_dice else (10, 5)
+
     plt.figure(figsize=figsize)
 
-    plt.subplot(1, 3, 1)
-    plt.plot(history["train_losses"], label="Train Loss")
-    plt.plot(history["val_losses"], label="Val Loss")
+    # Plot loss
+    plt.subplot(1, n_plots, 1)
+    if train_loss_key in history:
+        plt.plot(history[train_loss_key], label="Train Loss")
+    if val_loss_key in history:
+        plt.plot(history[val_loss_key], label="Val Loss")
     plt.title("Loss")
     plt.xlabel("Epoch")
     plt.ylabel("Loss")
     plt.legend()
     plt.grid(True)
 
-    plt.subplot(1, 3, 2)
-    plt.plot(history["val_ious"], label="Val IoU")
+    # Plot IoU
+    plt.subplot(1, n_plots, 2)
+    if val_iou_key in history:
+        plt.plot(history[val_iou_key], label="Val IoU")
     plt.title("IoU Score")
     plt.xlabel("Epoch")
     plt.ylabel("IoU")
     plt.legend()
     plt.grid(True)
 
-    plt.subplot(1, 3, 3)
-    plt.plot(history["val_dices"], label="Val Dice")
-    plt.title("Dice Score")
-    plt.xlabel("Epoch")
-    plt.ylabel("Dice")
-    plt.legend()
-    plt.grid(True)
+    # Plot Dice if available
+    if has_dice:
+        plt.subplot(1, n_plots, 3)
+        plt.plot(history[val_dice_key], label="Val Dice")
+        plt.title("Dice Score")
+        plt.xlabel("Epoch")
+        plt.ylabel("Dice")
+        plt.legend()
+        plt.grid(True)
 
     plt.tight_layout()
     plt.show()
 
     if verbose:
-        print(f"Best IoU: {max(history['val_ious']):.4f}")
-        print(f"Best Dice: {max(history['val_dices']):.4f}")
-        print(f"Final IoU: {history['val_ious'][-1]:.4f}")
-        print(f"Final Dice: {history['val_dices'][-1]:.4f}")
+        if val_iou_key in history:
+            print(f"Best IoU: {max(history[val_iou_key]):.4f}")
+            print(f"Final IoU: {history[val_iou_key][-1]:.4f}")
+        if val_dice_key in history:
+            print(f"Best Dice: {max(history[val_dice_key]):.4f}")
+            print(f"Final Dice: {history[val_dice_key][-1]:.4f}")
 
 
 def get_device():
@@ -7280,105 +7534,123 @@ def plot_prediction_comparison(
     prediction_colormap: str = "gray",
     ground_truth_colormap: str = "gray",
     original_colormap: Optional[str] = None,
+    indexes: Optional[List[int]] = None,
+    divider: Optional[float] = None,
 ):
-    """
-    Plot original image, prediction image, and optionally ground truth image side by side.
+    """Plot original image, prediction, and optional ground truth side by side.
+
+    Supports input as file paths, NumPy arrays, or PIL Images. For multi-band
+    images, selected channels can be specified via `indexes`. If the image data
+    is not normalized (e.g., Sentinel-2 [0, 10000]), the `divider` can be used
+    to scale values for visualization.
 
     Args:
-        original_image: Original input image (file path, numpy array, or PIL Image)
-        prediction_image: Prediction/segmentation mask (file path, numpy array, or PIL Image)
-        ground_truth_image: Optional ground truth mask (file path, numpy array, or PIL Image)
-        titles: Optional list of titles for each subplot
-        figsize: Figure size tuple (width, height)
-        save_path: Optional path to save the plot
-        show_plot: Whether to display the plot
-        prediction_colormap: Colormap for prediction image
-        ground_truth_colormap: Colormap for ground truth image
-        original_colormap: Colormap for original image (None for RGB)
+        original_image (Union[str, np.ndarray, Image.Image]):
+            Original input image as a file path, NumPy array, or PIL Image.
+        prediction_image (Union[str, np.ndarray, Image.Image]):
+            Predicted segmentation mask image.
+        ground_truth_image (Optional[Union[str, np.ndarray, Image.Image]], optional):
+            Ground truth mask image. Defaults to None.
+        titles (Optional[List[str]], optional):
+            List of titles for the subplots. If not provided, default titles are used.
+        figsize (Tuple[int, int], optional):
+            Size of the entire figure in inches. Defaults to (15, 5).
+        save_path (Optional[str], optional):
+            If specified, saves the figure to this path. Defaults to None.
+        show_plot (bool, optional):
+            Whether to display the figure using plt.show(). Defaults to True.
+        prediction_colormap (str, optional):
+            Colormap to use for the prediction mask. Defaults to "gray".
+        ground_truth_colormap (str, optional):
+            Colormap to use for the ground truth mask. Defaults to "gray".
+        original_colormap (Optional[str], optional):
+            Colormap to use for the original image if it's grayscale. Defaults to None.
+        indexes (Optional[List[int]], optional):
+            List of band/channel indexes (0-based for NumPy, 1-based for rasterio) to extract from the original image.
+            Useful for multi-band imagery like Sentinel-2. Defaults to None.
+        divider (Optional[float], optional):
+            Value to divide the original image by for normalization (e.g., 10000 for reflectance). Defaults to None.
 
     Returns:
-        matplotlib.figure.Figure: The figure object
+        matplotlib.figure.Figure:
+            The generated matplotlib figure object.
     """
 
-    def _load_image(img_input):
+    def _load_image(img_input, indexes=None):
         """Helper function to load image from various input types."""
         if isinstance(img_input, str):
-            # File path
             if img_input.lower().endswith((".tif", ".tiff")):
-                # Handle GeoTIFF files
                 with rasterio.open(img_input) as src:
-                    img = src.read()
-                    if img.shape[0] == 1:
-                        # Single band
-                        img = img[0]
+                    if indexes:
+                        img = src.read(indexes)  # 1-based
+                        img = (
+                            np.transpose(img, (1, 2, 0)) if len(indexes) > 1 else img[0]
+                        )
                     else:
-                        # Multi-band, transpose to (H, W, C)
-                        img = np.transpose(img, (1, 2, 0))
+                        img = src.read()
+                        if img.shape[0] == 1:
+                            img = img[0]
+                        else:
+                            img = np.transpose(img, (1, 2, 0))
             else:
-                # Regular image file
                 img = np.array(Image.open(img_input))
         elif isinstance(img_input, Image.Image):
-            # PIL Image
             img = np.array(img_input)
         elif isinstance(img_input, np.ndarray):
-            # NumPy array
             img = img_input
+            if indexes is not None and img.ndim == 3:
+                img = img[:, :, indexes]
         else:
             raise ValueError(f"Unsupported image type: {type(img_input)}")
-
         return img
 
     # Load images
-    original = _load_image(original_image)
+    original = _load_image(original_image, indexes=indexes)
     prediction = _load_image(prediction_image)
     ground_truth = (
         _load_image(ground_truth_image) if ground_truth_image is not None else None
     )
 
-    # Determine number of subplots
-    num_plots = 3 if ground_truth is not None else 2
+    # Apply divider normalization if requested
+    if divider is not None and isinstance(original, np.ndarray) and original.ndim == 3:
+        original = np.clip(original.astype(np.float32) / divider, 0, 1)
 
-    # Create figure and subplots
+    # Determine layout
+    num_plots = 3 if ground_truth is not None else 2
     fig, axes = plt.subplots(1, num_plots, figsize=figsize)
     if num_plots == 2:
         axes = [axes[0], axes[1]]
 
-    # Default titles
     if titles is None:
         titles = ["Original Image", "Prediction"]
         if ground_truth is not None:
             titles.append("Ground Truth")
 
-    # Plot original image
-    if len(original.shape) == 3 and original.shape[2] in [3, 4]:
-        # RGB or RGBA image
+    # Plot original
+    if original.ndim == 3 and original.shape[2] in [3, 4]:
         axes[0].imshow(original)
     else:
-        # Grayscale or single channel
         axes[0].imshow(original, cmap=original_colormap)
     axes[0].set_title(titles[0])
     axes[0].axis("off")
 
-    # Plot prediction image
+    # Prediction
     axes[1].imshow(prediction, cmap=prediction_colormap)
     axes[1].set_title(titles[1])
     axes[1].axis("off")
 
-    # Plot ground truth if provided
+    # Ground truth
     if ground_truth is not None:
         axes[2].imshow(ground_truth, cmap=ground_truth_colormap)
         axes[2].set_title(titles[2])
         axes[2].axis("off")
 
-    # Adjust layout
     plt.tight_layout()
 
-    # Save if requested
     if save_path:
         plt.savefig(save_path, dpi=300, bbox_inches="tight")
         print(f"Plot saved to: {save_path}")
 
-    # Show plot
     if show_plot:
         plt.show()
 

{geoai_py-0.8.1.dist-info → geoai_py-0.8.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geoai-py
-Version: 0.8.1
+Version: 0.8.3
 Summary: A Python package for using Artificial Intelligence (AI) with geospatial data
 Author-email: Qiusheng Wu <giswqs@gmail.com>
 License: MIT License

geoai_py-0.8.3.dist-info/RECORD

@@ -0,0 +1,17 @@
+geoai/__init__.py,sha256=Ks7SlSE7TkToYHB3EhXxIQLsYXuzfwHjyEUKJjkd8SA,3765
+geoai/classify.py,sha256=O8fah3DOBDMZW7V_qfDYsUnjB-9Wo5fjA-0e4wvUeAE,35054
+geoai/download.py,sha256=EQpcrcqMsYhDpd7bpjf4hGS5xL2oO-jsjngLgjjP3cE,46599
+geoai/extract.py,sha256=vyHH1k5zaXiy1SLdCLXxbWiNLp8XKdu_MXZoREMtAOQ,119102
+geoai/geoai.py,sha256=_Ar7PJgWpN86tm1YhzLVj7r1lYA3faRJT5brf1WzcwQ,9631
+geoai/hf.py,sha256=mLKGxEAS5eHkxZLwuLpYc1o7e3-7QIXdBv-QUY-RkFk,17072
+geoai/sam.py,sha256=O6S-kGiFn7YEcFbfWFItZZQOhnsm6-GlunxQLY0daEs,34345
+geoai/segment.py,sha256=pThAyq8kjgVDhMwHMiWkZ2qL5ykzA5lRg7tyMmSEBxk,43434
+geoai/segmentation.py,sha256=AtPzCvguHAEeuyXafa4bzMFATvltEYcah1B8ZMfkM_s,11373
+geoai/train.py,sha256=0uIz3i3sH-Knyq9wBrKLy22f--oLk2m-fCdhfwJosNU,129112
+geoai/utils.py,sha256=jNSXXQ056iyeDHDnkPUJ0_TLe-NonGbXh2YCN-gM50c,297270
+geoai_py-0.8.3.dist-info/licenses/LICENSE,sha256=vN2L5U7cZ6ZkOHFmc8WiGlsogWsZc5dllMeNxnKVOZg,1070
+geoai_py-0.8.3.dist-info/METADATA,sha256=ifNb98sm8g-Bsicc_R_9ljKahRXjibbEbSxxG8ga2mk,6661
+geoai_py-0.8.3.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+geoai_py-0.8.3.dist-info/entry_points.txt,sha256=uGp3Az3HURIsRHP9v-ys0hIbUuBBNUfXv6VbYHIXeg4,41
+geoai_py-0.8.3.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
+geoai_py-0.8.3.dist-info/RECORD,,

geoai_py-0.8.1.dist-info/RECORD

@@ -1,17 +0,0 @@
-geoai/__init__.py,sha256=NRaQhG5-sDLmpYXZXhbJ_LPkCEQbT-71qDAspoNEyxk,3765
-geoai/classify.py,sha256=O8fah3DOBDMZW7V_qfDYsUnjB-9Wo5fjA-0e4wvUeAE,35054
-geoai/download.py,sha256=EQpcrcqMsYhDpd7bpjf4hGS5xL2oO-jsjngLgjjP3cE,46599
-geoai/extract.py,sha256=vyHH1k5zaXiy1SLdCLXxbWiNLp8XKdu_MXZoREMtAOQ,119102
-geoai/geoai.py,sha256=HU6epCjpk228J65ZXmxY6GKlHg_ncmWV3UQUr_f8QTM,9447
-geoai/hf.py,sha256=mLKGxEAS5eHkxZLwuLpYc1o7e3-7QIXdBv-QUY-RkFk,17072
-geoai/sam.py,sha256=O6S-kGiFn7YEcFbfWFItZZQOhnsm6-GlunxQLY0daEs,34345
-geoai/segment.py,sha256=pThAyq8kjgVDhMwHMiWkZ2qL5ykzA5lRg7tyMmSEBxk,43434
-geoai/segmentation.py,sha256=AtPzCvguHAEeuyXafa4bzMFATvltEYcah1B8ZMfkM_s,11373
-geoai/train.py,sha256=Mrsb0yMVnprQHld3zDvA-puc-r8hGm1XgG0j2GGIn7E,112845
-geoai/utils.py,sha256=WQ_h39LpU65i-LP6GfFJLZBpiY0fyPdFc7dZcuyZ24s,284611
-geoai_py-0.8.1.dist-info/licenses/LICENSE,sha256=vN2L5U7cZ6ZkOHFmc8WiGlsogWsZc5dllMeNxnKVOZg,1070
-geoai_py-0.8.1.dist-info/METADATA,sha256=zOOCsS3BTaqIjl3DPsbxDDba4kK7kDKgNk2ytC_sYJ8,6661
-geoai_py-0.8.1.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-geoai_py-0.8.1.dist-info/entry_points.txt,sha256=uGp3Az3HURIsRHP9v-ys0hIbUuBBNUfXv6VbYHIXeg4,41
-geoai_py-0.8.1.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
-geoai_py-0.8.1.dist-info/RECORD,,