geoai-py 0.14.0__py2.py3-none-any.whl → 0.15.0__py2.py3-none-any.whl

geoai/utils.py

@@ -751,6 +751,12 @@ def view_vector_interactive(
         },
     }
 
+    # Make it compatible with binder and JupyterHub
+    if os.environ.get("JUPYTERHUB_SERVICE_PREFIX") is not None:
+        os.environ["LOCALTILESERVER_CLIENT_PREFIX"] = (
+            f"{os.environ['JUPYTERHUB_SERVICE_PREFIX'].lstrip('/')}/proxy/{{port}}"
+        )
+
     basemap_layer_name = None
     raster_layer = None
 
@@ -2609,6 +2615,7 @@ def export_geotiff_tiles(
     all_touched=True,
     create_overview=False,
     skip_empty_tiles=False,
+    metadata_format="PASCAL_VOC",
 ):
     """
     Export georeferenced GeoTIFF tiles and labels from raster and classification data.
@@ -2626,6 +2633,7 @@ def export_geotiff_tiles(
         all_touched (bool): Whether to use all_touched=True in rasterization (for vector data)
         create_overview (bool): Whether to create an overview image of all tiles
         skip_empty_tiles (bool): If True, skip tiles with no features
+        metadata_format (str): Output metadata format (PASCAL_VOC, COCO, YOLO). Default: PASCAL_VOC
     """
 
     import logging
@@ -2638,8 +2646,16 @@ def export_geotiff_tiles(
     os.makedirs(image_dir, exist_ok=True)
     label_dir = os.path.join(out_folder, "labels")
     os.makedirs(label_dir, exist_ok=True)
-    ann_dir = os.path.join(out_folder, "annotations")
-    os.makedirs(ann_dir, exist_ok=True)
+
+    # Create annotation directory based on metadata format
+    if metadata_format in ["PASCAL_VOC", "COCO"]:
+        ann_dir = os.path.join(out_folder, "annotations")
+        os.makedirs(ann_dir, exist_ok=True)
+
+    # Initialize COCO annotations dictionary
+    if metadata_format == "COCO":
+        coco_annotations = {"images": [], "annotations": [], "categories": []}
+        ann_id = 0
 
     # Determine if class data is raster or vector
     is_class_data_raster = False
@@ -2713,6 +2729,17 @@ def export_geotiff_tiles(
 
                 # Create class mapping
                 class_to_id = {int(cls): i + 1 for i, cls in enumerate(unique_classes)}
+
+                # Populate COCO categories
+                if metadata_format == "COCO":
+                    for cls_val in unique_classes:
+                        coco_annotations["categories"].append(
+                            {
+                                "id": class_to_id[int(cls_val)],
+                                "name": str(int(cls_val)),
+                                "supercategory": "object",
+                            }
+                        )
         else:
             # Load vector class data
             try:
@@ -2742,12 +2769,33 @@ def export_geotiff_tiles(
                         )
                     # Create class mapping
                     class_to_id = {cls: i + 1 for i, cls in enumerate(unique_classes)}
+
+                    # Populate COCO categories
+                    if metadata_format == "COCO":
+                        for cls_val in unique_classes:
+                            coco_annotations["categories"].append(
+                                {
+                                    "id": class_to_id[cls_val],
+                                    "name": str(cls_val),
+                                    "supercategory": "object",
+                                }
+                            )
                 else:
                     if not quiet:
                         print(
                             f"WARNING: '{class_value_field}' not found in vector data. Using default class ID 1."
                         )
                     class_to_id = {1: 1}  # Default mapping
+
+                    # Populate COCO categories with default
+                    if metadata_format == "COCO":
+                        coco_annotations["categories"].append(
+                            {
+                                "id": 1,
+                                "name": "object",
+                                "supercategory": "object",
+                            }
+                        )
             except Exception as e:
                 raise ValueError(f"Error processing vector data: {e}")
 
@@ -2964,72 +3012,186 @@ def export_geotiff_tiles(
                     pbar.write(f"ERROR saving label GeoTIFF: {e}")
                     stats["errors"] += 1
 
-                # Create XML annotation for object detection if using vector class data
+                # Create annotations for object detection if using vector class data
                 if (
                     not is_class_data_raster
                     and "gdf" in locals()
                     and len(window_features) > 0
                 ):
-                    # Create XML annotation
-                    root = ET.Element("annotation")
-                    ET.SubElement(root, "folder").text = "images"
-                    ET.SubElement(root, "filename").text = f"tile_{tile_index:06d}.tif"
+                    if metadata_format == "PASCAL_VOC":
+                        # Create XML annotation
+                        root = ET.Element("annotation")
+                        ET.SubElement(root, "folder").text = "images"
+                        ET.SubElement(root, "filename").text = (
+                            f"tile_{tile_index:06d}.tif"
+                        )
 
-                    size = ET.SubElement(root, "size")
-                    ET.SubElement(size, "width").text = str(tile_size)
-                    ET.SubElement(size, "height").text = str(tile_size)
-                    ET.SubElement(size, "depth").text = str(image_data.shape[0])
+                        size = ET.SubElement(root, "size")
+                        ET.SubElement(size, "width").text = str(tile_size)
+                        ET.SubElement(size, "height").text = str(tile_size)
+                        ET.SubElement(size, "depth").text = str(image_data.shape[0])
+
+                        # Add georeference information
+                        geo = ET.SubElement(root, "georeference")
+                        ET.SubElement(geo, "crs").text = str(src.crs)
+                        ET.SubElement(geo, "transform").text = str(
+                            window_transform
+                        ).replace("\n", "")
+                        ET.SubElement(geo, "bounds").text = (
+                            f"{minx}, {miny}, {maxx}, {maxy}"
+                        )
 
-                    # Add georeference information
-                    geo = ET.SubElement(root, "georeference")
-                    ET.SubElement(geo, "crs").text = str(src.crs)
-                    ET.SubElement(geo, "transform").text = str(
-                        window_transform
-                    ).replace("\n", "")
-                    ET.SubElement(geo, "bounds").text = (
-                        f"{minx}, {miny}, {maxx}, {maxy}"
-                    )
+                        # Add objects
+                        for idx, feature in window_features.iterrows():
+                            # Get feature class
+                            if class_value_field in feature:
+                                class_val = feature[class_value_field]
+                            else:
+                                class_val = "object"
 
-                    # Add objects
-                    for idx, feature in window_features.iterrows():
-                        # Get feature class
-                        if class_value_field in feature:
-                            class_val = feature[class_value_field]
-                        else:
-                            class_val = "object"
+                            # Get geometry bounds in pixel coordinates
+                            geom = feature.geometry.intersection(window_bounds)
+                            if not geom.is_empty:
+                                # Get bounds in world coordinates
+                                minx_f, miny_f, maxx_f, maxy_f = geom.bounds
+
+                                # Convert to pixel coordinates
+                                col_min, row_min = ~window_transform * (minx_f, maxy_f)
+                                col_max, row_max = ~window_transform * (maxx_f, miny_f)
+
+                                # Ensure coordinates are within tile bounds
+                                xmin = max(0, min(tile_size, int(col_min)))
+                                ymin = max(0, min(tile_size, int(row_min)))
+                                xmax = max(0, min(tile_size, int(col_max)))
+                                ymax = max(0, min(tile_size, int(row_max)))
+
+                                # Only add if the box has non-zero area
+                                if xmax > xmin and ymax > ymin:
+                                    obj = ET.SubElement(root, "object")
+                                    ET.SubElement(obj, "name").text = str(class_val)
+                                    ET.SubElement(obj, "difficult").text = "0"
+
+                                    bbox = ET.SubElement(obj, "bndbox")
+                                    ET.SubElement(bbox, "xmin").text = str(xmin)
+                                    ET.SubElement(bbox, "ymin").text = str(ymin)
+                                    ET.SubElement(bbox, "xmax").text = str(xmax)
+                                    ET.SubElement(bbox, "ymax").text = str(ymax)
+
+                        # Save XML
+                        tree = ET.ElementTree(root)
+                        xml_path = os.path.join(ann_dir, f"tile_{tile_index:06d}.xml")
+                        tree.write(xml_path)
 
-                        # Get geometry bounds in pixel coordinates
-                        geom = feature.geometry.intersection(window_bounds)
-                        if not geom.is_empty:
-                            # Get bounds in world coordinates
-                            minx_f, miny_f, maxx_f, maxy_f = geom.bounds
-
-                            # Convert to pixel coordinates
-                            col_min, row_min = ~window_transform * (minx_f, maxy_f)
-                            col_max, row_max = ~window_transform * (maxx_f, miny_f)
-
-                            # Ensure coordinates are within tile bounds
-                            xmin = max(0, min(tile_size, int(col_min)))
-                            ymin = max(0, min(tile_size, int(row_min)))
-                            xmax = max(0, min(tile_size, int(col_max)))
-                            ymax = max(0, min(tile_size, int(row_max)))
-
-                            # Only add if the box has non-zero area
-                            if xmax > xmin and ymax > ymin:
-                                obj = ET.SubElement(root, "object")
-                                ET.SubElement(obj, "name").text = str(class_val)
-                                ET.SubElement(obj, "difficult").text = "0"
-
-                                bbox = ET.SubElement(obj, "bndbox")
-                                ET.SubElement(bbox, "xmin").text = str(xmin)
-                                ET.SubElement(bbox, "ymin").text = str(ymin)
-                                ET.SubElement(bbox, "xmax").text = str(xmax)
-                                ET.SubElement(bbox, "ymax").text = str(ymax)
+                    elif metadata_format == "COCO":
+                        # Add image info
+                        image_id = tile_index
+                        coco_annotations["images"].append(
+                            {
+                                "id": image_id,
+                                "file_name": f"tile_{tile_index:06d}.tif",
+                                "width": tile_size,
+                                "height": tile_size,
+                                "crs": str(src.crs),
+                                "transform": str(window_transform),
+                            }
+                        )
 
-                    # Save XML
-                    tree = ET.ElementTree(root)
-                    xml_path = os.path.join(ann_dir, f"tile_{tile_index:06d}.xml")
-                    tree.write(xml_path)
+                        # Add annotations for each feature
+                        for _, feature in window_features.iterrows():
+                            # Get feature class
+                            if class_value_field in feature:
+                                class_val = feature[class_value_field]
+                                category_id = class_to_id.get(class_val, 1)
+                            else:
+                                category_id = 1
+
+                            # Get geometry bounds
+                            geom = feature.geometry.intersection(window_bounds)
+                            if not geom.is_empty:
+                                # Get bounds in world coordinates
+                                minx_f, miny_f, maxx_f, maxy_f = geom.bounds
+
+                                # Convert to pixel coordinates
+                                col_min, row_min = ~window_transform * (minx_f, maxy_f)
+                                col_max, row_max = ~window_transform * (maxx_f, miny_f)
+
+                                # Ensure coordinates are within tile bounds
+                                xmin = max(0, min(tile_size, int(col_min)))
+                                ymin = max(0, min(tile_size, int(row_min)))
+                                xmax = max(0, min(tile_size, int(col_max)))
+                                ymax = max(0, min(tile_size, int(row_max)))
+
+                                # Skip if box is too small
+                                if xmax - xmin < 1 or ymax - ymin < 1:
+                                    continue
+
+                                width = xmax - xmin
+                                height = ymax - ymin
+
+                                # Add annotation
+                                ann_id += 1
+                                coco_annotations["annotations"].append(
+                                    {
+                                        "id": ann_id,
+                                        "image_id": image_id,
+                                        "category_id": category_id,
+                                        "bbox": [xmin, ymin, width, height],
+                                        "area": width * height,
+                                        "iscrowd": 0,
+                                    }
+                                )
+
+                    elif metadata_format == "YOLO":
+                        # Create YOLO format annotations
+                        yolo_annotations = []
+
+                        for _, feature in window_features.iterrows():
+                            # Get feature class
+                            if class_value_field in feature:
+                                class_val = feature[class_value_field]
+                                # YOLO uses 0-indexed class IDs
+                                class_id = class_to_id.get(class_val, 1) - 1
+                            else:
+                                class_id = 0
+
+                            # Get geometry bounds
+                            geom = feature.geometry.intersection(window_bounds)
+                            if not geom.is_empty:
+                                # Get bounds in world coordinates
+                                minx_f, miny_f, maxx_f, maxy_f = geom.bounds
+
+                                # Convert to pixel coordinates
+                                col_min, row_min = ~window_transform * (minx_f, maxy_f)
+                                col_max, row_max = ~window_transform * (maxx_f, miny_f)
+
+                                # Ensure coordinates are within tile bounds
+                                xmin = max(0, min(tile_size, col_min))
+                                ymin = max(0, min(tile_size, row_min))
+                                xmax = max(0, min(tile_size, col_max))
+                                ymax = max(0, min(tile_size, row_max))
+
+                                # Skip if box is too small
+                                if xmax - xmin < 1 or ymax - ymin < 1:
+                                    continue
+
+                                # Calculate normalized coordinates (YOLO format)
+                                x_center = ((xmin + xmax) / 2) / tile_size
+                                y_center = ((ymin + ymax) / 2) / tile_size
+                                width = (xmax - xmin) / tile_size
+                                height = (ymax - ymin) / tile_size
+
+                                # Add YOLO annotation line
+                                yolo_annotations.append(
+                                    f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
+                                )
+
+                        # Save YOLO annotations to text file
+                        if yolo_annotations:
+                            yolo_path = os.path.join(
+                                label_dir, f"tile_{tile_index:06d}.txt"
+                            )
+                            with open(yolo_path, "w") as f:
+                                f.write("\n".join(yolo_annotations))
 
                 # Update progress bar
                 pbar.update(1)
@@ -3047,6 +3209,39 @@ def export_geotiff_tiles(
         # Close progress bar
         pbar.close()
 
+        # Save COCO annotations if applicable
+        if metadata_format == "COCO":
+            try:
+                with open(os.path.join(ann_dir, "instances.json"), "w") as f:
+                    json.dump(coco_annotations, f, indent=2)
+                if not quiet:
+                    print(
+                        f"Saved COCO annotations: {len(coco_annotations['images'])} images, "
+                        f"{len(coco_annotations['annotations'])} annotations, "
+                        f"{len(coco_annotations['categories'])} categories"
+                    )
+            except Exception as e:
+                if not quiet:
+                    print(f"ERROR saving COCO annotations: {e}")
+                stats["errors"] += 1
+
+        # Save YOLO classes file if applicable
+        if metadata_format == "YOLO":
+            try:
+                # Create classes.txt with class names
+                classes_path = os.path.join(out_folder, "classes.txt")
+                # Sort by class ID to ensure correct order
+                sorted_classes = sorted(class_to_id.items(), key=lambda x: x[1])
+                with open(classes_path, "w") as f:
+                    for class_val, _ in sorted_classes:
+                        f.write(f"{class_val}\n")
+                if not quiet:
+                    print(f"Saved YOLO classes file with {len(class_to_id)} classes")
+            except Exception as e:
+                if not quiet:
+                    print(f"ERROR saving YOLO classes file: {e}")
+                stats["errors"] += 1
+
         # Create overview image if requested
         if create_overview and stats["tile_coordinates"]:
             try:
@@ -3125,11 +3320,11 @@ def export_geotiff_tiles_batch(
     max_tiles=None,
     quiet=False,
     all_touched=True,
-    create_overview=False,
     skip_empty_tiles=False,
     image_extensions=None,
     mask_extensions=None,
     match_by_name=True,
+    metadata_format="PASCAL_VOC",
 ) -> Dict[str, Any]:
     """
     Export georeferenced GeoTIFF tiles from images and masks.
@@ -3170,6 +3365,8 @@ def export_geotiff_tiles_batch(
         mask_extensions (list): List of mask file extensions to process (default: common raster/vector formats)
         match_by_name (bool): If True, match image and mask files by base filename.
             If False, match by sorted order (alphabetically). Only applies when masks_folder is used.
+        metadata_format (str): Annotation format - "PASCAL_VOC" (XML), "COCO" (JSON), or "YOLO" (TXT).
+            Default is "PASCAL_VOC".
 
     Returns:
         Dict[str, Any]: Dictionary containing batch processing statistics
@@ -3253,6 +3450,19 @@ def export_geotiff_tiles_batch(
     os.makedirs(output_images_dir, exist_ok=True)
     os.makedirs(output_masks_dir, exist_ok=True)
 
+    # Create annotation directory based on metadata format
+    if metadata_format in ["PASCAL_VOC", "COCO"]:
+        ann_dir = os.path.join(output_folder, "annotations")
+        os.makedirs(ann_dir, exist_ok=True)
+
+    # Initialize COCO annotations dictionary
+    coco_annotations = None
+    if metadata_format == "COCO":
+        coco_annotations = {"images": [], "annotations": [], "categories": []}
+
+    # Initialize YOLO class set
+    yolo_classes = set() if metadata_format == "YOLO" else None
+
     # Get list of image files
     image_files = []
     for ext in image_extensions:
@@ -3406,6 +3616,13 @@ def export_geotiff_tiles_batch(
                 quiet=quiet,
                 mask_gdf=mask_gdf,  # Pass pre-loaded GeoDataFrame if using single mask
                 use_single_mask_file=use_single_mask_file,
+                metadata_format=metadata_format,
+                ann_dir=(
+                    ann_dir
+                    if "ann_dir" in locals()
+                    and metadata_format in ["PASCAL_VOC", "COCO"]
+                    else None
+                ),
             )
 
             # Update counters
@@ -3427,6 +3644,23 @@ def export_geotiff_tiles_batch(
                 }
             )
 
+            # Aggregate COCO annotations
+            if metadata_format == "COCO" and "coco_data" in tiles_generated:
+                coco_data = tiles_generated["coco_data"]
+                # Add images and annotations
+                coco_annotations["images"].extend(coco_data.get("images", []))
+                coco_annotations["annotations"].extend(coco_data.get("annotations", []))
+                # Merge categories (avoid duplicates)
+                for cat in coco_data.get("categories", []):
+                    if not any(
+                        c["id"] == cat["id"] for c in coco_annotations["categories"]
+                    ):
+                        coco_annotations["categories"].append(cat)
+
+            # Aggregate YOLO classes
+            if metadata_format == "YOLO" and "yolo_classes" in tiles_generated:
+                yolo_classes.update(tiles_generated["yolo_classes"])
+
         except Exception as e:
             if not quiet:
                 print(f"ERROR processing {base_name}: {e}")
@@ -3435,6 +3669,33 @@ def export_geotiff_tiles_batch(
             )
             batch_stats["errors"] += 1
 
+    # Save aggregated COCO annotations
+    if metadata_format == "COCO" and coco_annotations:
+        import json
+
+        coco_path = os.path.join(ann_dir, "instances.json")
+        with open(coco_path, "w") as f:
+            json.dump(coco_annotations, f, indent=2)
+        if not quiet:
+            print(f"\nSaved COCO annotations: {coco_path}")
+            print(
+                f"  Images: {len(coco_annotations['images'])}, "
+                f"Annotations: {len(coco_annotations['annotations'])}, "
+                f"Categories: {len(coco_annotations['categories'])}"
+            )
+
+    # Save aggregated YOLO classes
+    if metadata_format == "YOLO" and yolo_classes:
+        classes_path = os.path.join(output_folder, "labels", "classes.txt")
+        os.makedirs(os.path.dirname(classes_path), exist_ok=True)
+        sorted_classes = sorted(yolo_classes)
+        with open(classes_path, "w") as f:
+            for cls in sorted_classes:
+                f.write(f"{cls}\n")
+        if not quiet:
+            print(f"\nSaved YOLO classes: {classes_path}")
+            print(f"  Total classes: {len(sorted_classes)}")
+
     # Print batch summary
     if not quiet:
         print("\n" + "=" * 60)
@@ -3458,6 +3719,10 @@ def export_geotiff_tiles_batch(
         print(f"Output saved to: {output_folder}")
         print(f"  Images: {output_images_dir}")
         print(f"  Masks: {output_masks_dir}")
+        if metadata_format in ["PASCAL_VOC", "COCO"]:
+            print(f"  Annotations: {ann_dir}")
+        elif metadata_format == "YOLO":
+            print(f"  Labels: {os.path.join(output_folder, 'labels')}")
 
         # List failed files if any
         if batch_stats["failed_files"]:
@@ -3485,6 +3750,8 @@ def _process_image_mask_pair(
     quiet=False,
     mask_gdf=None,
     use_single_mask_file=False,
+    metadata_format="PASCAL_VOC",
+    ann_dir=None,
 ):
     """
     Process a single image-mask pair and save tiles directly to output directories.
@@ -3517,6 +3784,13 @@ def _process_image_mask_pair(
         "errors": 0,
     }
 
+    # Initialize COCO/YOLO tracking for this image
+    if metadata_format == "COCO":
+        stats["coco_data"] = {"images": [], "annotations": [], "categories": []}
+        coco_ann_id = 0
+    if metadata_format == "YOLO":
+        stats["yolo_classes"] = set()
+
     # Open the input raster
     with rasterio.open(image_file) as src:
         # Calculate number of tiles
@@ -3609,9 +3883,6 @@ def _process_image_mask_pair(
         tile_index = 0
         for y in range(num_tiles_y):
             for x in range(num_tiles_x):
-                if tile_index >= max_tiles:
-                    break
-
                 # Calculate window coordinates
                 window_x = x * stride
                 window_y = y * stride
@@ -3714,9 +3985,12 @@ def _process_image_mask_pair(
 
                 # Skip tile if no features and skip_empty_tiles is True
                 if skip_empty_tiles and not has_features:
-                    tile_index += 1
                     continue
 
+                # Check if we've reached max_tiles before saving
+                if tile_index >= max_tiles:
+                    break
+
                 # Generate unique tile name
                 tile_name = f"{base_name}_{global_tile_counter + tile_index:06d}"
 
@@ -3771,6 +4045,197 @@ def _process_image_mask_pair(
                         print(f"ERROR saving label GeoTIFF: {e}")
                     stats["errors"] += 1
 
+                # Generate annotation metadata based on format
+                if metadata_format == "PASCAL_VOC" and ann_dir:
+                    # Create PASCAL VOC XML annotation
+                    from lxml import etree as ET
+
+                    annotation = ET.Element("annotation")
+                    ET.SubElement(annotation, "folder").text = os.path.basename(
+                        output_images_dir
+                    )
+                    ET.SubElement(annotation, "filename").text = f"{tile_name}.tif"
+                    ET.SubElement(annotation, "path").text = image_path
+
+                    source = ET.SubElement(annotation, "source")
+                    ET.SubElement(source, "database").text = "GeoAI"
+
+                    size = ET.SubElement(annotation, "size")
+                    ET.SubElement(size, "width").text = str(tile_size)
+                    ET.SubElement(size, "height").text = str(tile_size)
+                    ET.SubElement(size, "depth").text = str(image_data.shape[0])
+
+                    ET.SubElement(annotation, "segmented").text = "1"
+
+                    # Find connected components for instance segmentation
+                    from scipy import ndimage
+
+                    for class_id in np.unique(label_mask):
+                        if class_id == 0:
+                            continue
+
+                        class_mask = (label_mask == class_id).astype(np.uint8)
+                        labeled_array, num_features = ndimage.label(class_mask)
+
+                        for instance_id in range(1, num_features + 1):
+                            instance_mask = labeled_array == instance_id
+                            coords = np.argwhere(instance_mask)
+
+                            if len(coords) == 0:
+                                continue
+
+                            ymin, xmin = coords.min(axis=0)
+                            ymax, xmax = coords.max(axis=0)
+
+                            obj = ET.SubElement(annotation, "object")
+                            class_name = next(
+                                (k for k, v in class_to_id.items() if v == class_id),
+                                str(class_id),
+                            )
+                            ET.SubElement(obj, "name").text = str(class_name)
+                            ET.SubElement(obj, "pose").text = "Unspecified"
+                            ET.SubElement(obj, "truncated").text = "0"
+                            ET.SubElement(obj, "difficult").text = "0"
+
+                            bndbox = ET.SubElement(obj, "bndbox")
+                            ET.SubElement(bndbox, "xmin").text = str(int(xmin))
+                            ET.SubElement(bndbox, "ymin").text = str(int(ymin))
+                            ET.SubElement(bndbox, "xmax").text = str(int(xmax))
+                            ET.SubElement(bndbox, "ymax").text = str(int(ymax))
+
+                    # Save XML file
+                    xml_path = os.path.join(ann_dir, f"{tile_name}.xml")
+                    tree = ET.ElementTree(annotation)
+                    tree.write(xml_path, pretty_print=True, encoding="utf-8")
+
+                elif metadata_format == "COCO":
+                    # Add COCO image entry
+                    image_id = int(global_tile_counter + tile_index)
+                    stats["coco_data"]["images"].append(
+                        {
+                            "id": image_id,
+                            "file_name": f"{tile_name}.tif",
+                            "width": int(tile_size),
+                            "height": int(tile_size),
+                        }
+                    )
+
+                    # Add COCO categories (only once per unique class)
+                    for class_val, class_id in class_to_id.items():
+                        if not any(
+                            c["id"] == class_id
+                            for c in stats["coco_data"]["categories"]
+                        ):
+                            stats["coco_data"]["categories"].append(
+                                {
+                                    "id": int(class_id),
+                                    "name": str(class_val),
+                                    "supercategory": "object",
+                                }
+                            )
+
+                    # Add COCO annotations (instance segmentation)
+                    from scipy import ndimage
+                    from skimage import measure
+
+                    for class_id in np.unique(label_mask):
+                        if class_id == 0:
+                            continue
+
+                        class_mask = (label_mask == class_id).astype(np.uint8)
+                        labeled_array, num_features = ndimage.label(class_mask)
+
+                        for instance_id in range(1, num_features + 1):
+                            instance_mask = (labeled_array == instance_id).astype(
+                                np.uint8
+                            )
+                            coords = np.argwhere(instance_mask)
+
+                            if len(coords) == 0:
+                                continue
+
+                            ymin, xmin = coords.min(axis=0)
+                            ymax, xmax = coords.max(axis=0)
+
+                            bbox = [
+                                int(xmin),
+                                int(ymin),
+                                int(xmax - xmin),
+                                int(ymax - ymin),
+                            ]
+                            area = int(np.sum(instance_mask))
+
+                            # Find contours for segmentation
+                            contours = measure.find_contours(instance_mask, 0.5)
+                            segmentation = []
+                            for contour in contours:
+                                contour = np.flip(contour, axis=1)
+                                segmentation_points = contour.ravel().tolist()
+                                if len(segmentation_points) >= 6:
+                                    segmentation.append(segmentation_points)
+
+                            if segmentation:
+                                stats["coco_data"]["annotations"].append(
+                                    {
+                                        "id": int(coco_ann_id),
+                                        "image_id": int(image_id),
+                                        "category_id": int(class_id),
+                                        "bbox": bbox,
+                                        "area": area,
+                                        "segmentation": segmentation,
+                                        "iscrowd": 0,
+                                    }
+                                )
+                                coco_ann_id += 1
+
+                elif metadata_format == "YOLO":
+                    # Create YOLO labels directory if needed
+                    labels_dir = os.path.join(
+                        os.path.dirname(output_images_dir), "labels"
+                    )
+                    os.makedirs(labels_dir, exist_ok=True)
+
+                    # Generate YOLO annotation file
+                    yolo_path = os.path.join(labels_dir, f"{tile_name}.txt")
+                    from scipy import ndimage
+
+                    with open(yolo_path, "w") as yolo_file:
+                        for class_id in np.unique(label_mask):
+                            if class_id == 0:
+                                continue
+
+                            # Track class for classes.txt
+                            class_name = next(
+                                (k for k, v in class_to_id.items() if v == class_id),
+                                str(class_id),
+                            )
+                            stats["yolo_classes"].add(class_name)
+
+                            class_mask = (label_mask == class_id).astype(np.uint8)
+                            labeled_array, num_features = ndimage.label(class_mask)
+
+                            for instance_id in range(1, num_features + 1):
+                                instance_mask = labeled_array == instance_id
+                                coords = np.argwhere(instance_mask)
+
+                                if len(coords) == 0:
+                                    continue
+
+                                ymin, xmin = coords.min(axis=0)
+                                ymax, xmax = coords.max(axis=0)
+
+                                # Convert to YOLO format (normalized center coordinates)
+                                x_center = ((xmin + xmax) / 2) / tile_size
+                                y_center = ((ymin + ymax) / 2) / tile_size
+                                width = (xmax - xmin) / tile_size
+                                height = (ymax - ymin) / tile_size
+
+                                # YOLO uses 0-based class indices
+                                yolo_class_id = class_id - 1
+                                yolo_file.write(
+                                    f"{yolo_class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n"
+                                )
+
                 tile_index += 1
                 if tile_index >= max_tiles:
                     break
@@ -3781,6 +4246,179 @@ def _process_image_mask_pair(
     return stats
 
 
+def display_training_tiles(
+    output_dir,
+    num_tiles=6,
+    figsize=(18, 6),
+    cmap="gray",
+    save_path=None,
+):
+    """
+    Display image and mask tile pairs from training data output.
+
+    Args:
+        output_dir (str): Path to output directory containing 'images' and 'masks' subdirectories
+        num_tiles (int): Number of tile pairs to display (default: 6)
+        figsize (tuple): Figure size as (width, height) in inches (default: (18, 6))
+        cmap (str): Colormap for mask display (default: 'gray')
+        save_path (str, optional): If provided, save figure to this path instead of displaying
+
+    Returns:
+        tuple: (fig, axes) matplotlib figure and axes objects
+
+    Example:
+        >>> fig, axes = display_training_tiles('output/tiles', num_tiles=6)
+        >>> # Or save to file
+        >>> display_training_tiles('output/tiles', num_tiles=4, save_path='tiles_preview.png')
+    """
+    import matplotlib.pyplot as plt
+
+    # Get list of image tiles
+    images_dir = os.path.join(output_dir, "images")
+    if not os.path.exists(images_dir):
+        raise ValueError(f"Images directory not found: {images_dir}")
+
+    image_tiles = sorted(os.listdir(images_dir))[:num_tiles]
+
+    if not image_tiles:
+        raise ValueError(f"No image tiles found in {images_dir}")
+
+    # Limit to available tiles
+    num_tiles = min(num_tiles, len(image_tiles))
+
+    # Create figure with subplots
+    fig, axes = plt.subplots(2, num_tiles, figsize=figsize)
+
+    # Handle case where num_tiles is 1
+    if num_tiles == 1:
+        axes = axes.reshape(2, 1)
+
+    for idx, tile_name in enumerate(image_tiles):
+        # Load and display image tile
+        image_path = os.path.join(output_dir, "images", tile_name)
+        with rasterio.open(image_path) as src:
+            show(src, ax=axes[0, idx], title=f"Image {idx+1}")
+
+        # Load and display mask tile
+        mask_path = os.path.join(output_dir, "masks", tile_name)
+        if os.path.exists(mask_path):
+            with rasterio.open(mask_path) as src:
+                show(src, ax=axes[1, idx], title=f"Mask {idx+1}", cmap=cmap)
+        else:
+            axes[1, idx].text(
+                0.5,
+                0.5,
+                "Mask not found",
+                ha="center",
+                va="center",
+                transform=axes[1, idx].transAxes,
+            )
+            axes[1, idx].set_title(f"Mask {idx+1}")
+
+    plt.tight_layout()
+
+    # Save or show
+    if save_path:
+        plt.savefig(save_path, dpi=150, bbox_inches="tight")
+        plt.close(fig)
+        print(f"Figure saved to: {save_path}")
+    else:
+        plt.show()
+
+    return fig, axes
+
+
+def display_image_with_vector(
+    image_path,
+    vector_path,
+    figsize=(16, 8),
+    vector_color="red",
+    vector_linewidth=1,
+    vector_facecolor="none",
+    save_path=None,
+):
+    """
+    Display a raster image alongside the same image with vector overlay.
+
+    Args:
+        image_path (str): Path to raster image file
+        vector_path (str): Path to vector file (GeoJSON, Shapefile, etc.)
+        figsize (tuple): Figure size as (width, height) in inches (default: (16, 8))
+        vector_color (str): Edge color for vector features (default: 'red')
+        vector_linewidth (float): Line width for vector features (default: 1)
+        vector_facecolor (str): Fill color for vector features (default: 'none')
+        save_path (str, optional): If provided, save figure to this path instead of displaying
+
+    Returns:
+        tuple: (fig, axes, info_dict) where info_dict contains image and vector metadata
+
+    Example:
+        >>> fig, axes, info = display_image_with_vector(
+        ...     'image.tif',
+        ...     'buildings.geojson',
+        ...     vector_color='blue'
+        ... )
+        >>> print(f"Number of features: {info['num_features']}")
+    """
+    import matplotlib.pyplot as plt
+
+    # Validate inputs
+    if not os.path.exists(image_path):
+        raise ValueError(f"Image file not found: {image_path}")
+    if not os.path.exists(vector_path):
+        raise ValueError(f"Vector file not found: {vector_path}")
+
+    # Create figure
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
+
+    # Load and display image
+    with rasterio.open(image_path) as src:
+        # Plot image only
+        show(src, ax=ax1, title="Image")
+
+        # Load vector data
+        vector_data = gpd.read_file(vector_path)
+
+        # Reproject to image CRS if needed
+        if vector_data.crs != src.crs:
+            vector_data = vector_data.to_crs(src.crs)
+
+        # Plot image with vector overlay
+        show(
+            src,
+            ax=ax2,
+            title=f"Image with {len(vector_data)} Vector Features",
+        )
+        vector_data.plot(
+            ax=ax2,
+            facecolor=vector_facecolor,
+            edgecolor=vector_color,
+            linewidth=vector_linewidth,
+        )
+
+        # Collect metadata
+        info = {
+            "image_shape": src.shape,
+            "image_crs": src.crs,
+            "image_bounds": src.bounds,
+            "num_features": len(vector_data),
+            "vector_crs": vector_data.crs,
+            "vector_bounds": vector_data.total_bounds,
+        }
+
+    plt.tight_layout()
+
+    # Save or show
+    if save_path:
+        plt.savefig(save_path, dpi=150, bbox_inches="tight")
+        plt.close(fig)
+        print(f"Figure saved to: {save_path}")
+    else:
+        plt.show()
+
+    return fig, (ax1, ax2), info
+
+
 def create_overview_image(
     src, tile_coordinates, output_path, tile_size, stride, geojson_path=None
 ) -> str: