geoai-py 0.8.2__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.2"
+__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():

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geoai-py
-Version: 0.8.2
+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.2.dist-info/RECORD

@@ -1,17 +0,0 @@
-geoai/__init__.py,sha256=wKdVPgcUCURuMI3TNzqh1MPw96Bs5XL7mNUST75bOjw,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=_JVEhFUzOdDS_Rmco8c4BJVnrs2VvUVHKf4ubtim5fg,286109
-geoai_py-0.8.2.dist-info/licenses/LICENSE,sha256=vN2L5U7cZ6ZkOHFmc8WiGlsogWsZc5dllMeNxnKVOZg,1070
-geoai_py-0.8.2.dist-info/METADATA,sha256=6LIaLEjT5AtbV5OOJwMT7NQCMI6pHcKXsfQbakWl6Ns,6661
-geoai_py-0.8.2.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-geoai_py-0.8.2.dist-info/entry_points.txt,sha256=uGp3Az3HURIsRHP9v-ys0hIbUuBBNUfXv6VbYHIXeg4,41
-geoai_py-0.8.2.dist-info/top_level.txt,sha256=1YkCUWu-ii-0qIex7kbwAvfei-gos9ycyDyUCJPNWHY,6
-geoai_py-0.8.2.dist-info/RECORD,,