readtopos.py

import marimo

__generated_with = "0.20.1"
app = marimo.App(width="medium")


@app.cell
def loadskill():
    with open("skills.md", "r") as f:
        print(f.read())
    return


@app.cell
def _():
    # load libraries and the folder of topos
    import marimo as mo
    import pathlib
    import rioxarray
    import xarray as xr
    import folium

    # Define the path to your tiff directory
    tiff_dir = pathlib.Path("tiffs")

    # Get a list of all .tif and .tiff files
    tiff_files = sorted([f.name for f in tiff_dir.glob("*.tif*")])

    # Create a selection UI
    #tiff_selector = mo.ui.dropdown(
    #    options=tiff_files, 
    #    label="Select Soviet Map Sheet (TIFF)",
    #    value=tiff_files[0] if tiff_files else None
    #)

    mo.md(
        f"""
        # Maps Explorer
        Found **{len(tiff_files)}** files in `{tiff_dir}/`.


        """
    )
    return folium, mo, pathlib, rioxarray, tiff_dir, tiff_files


@app.cell
def _(mo, tiff_files):
    # this is a setup cell
    # choose your .tiff and set your parameters.
    import numpy as np
    import polars as pl
    import geopandas as gpd
    from skimage import measure, morphology
    import pyproj
    import shapely.geometry as sg
    import os

    # Setup UI Elements as global variables so their values carry through the app
    tiff_selector = mo.ui.dropdown(
        options=tiff_files, 
        label="Select Soviet Map Sheet", 
        value=tiff_files[0] if tiff_files else None
    )
    threshold_slider = mo.ui.slider(
        0, 255, 
        value=90, 
        label="Darkness Threshold"
    )
    min_building_size = mo.ui.number(
        start=10, 
        stop=2000, 
        value=150, 
        label="Min Building Size (px)"
    )

    # Display the UI
    mo.vstack([
         mo.md("# Soviet Map Vectorizer"),
         mo.hstack([tiff_selector, threshold_slider, min_building_size])
    ])
    return (
        gpd,
        measure,
        min_building_size,
        morphology,
        np,
        pl,
        pyproj,
        sg,
        threshold_slider,
        tiff_selector,
    )


@app.cell
def _(tiff_selector):
    def _():
        # Load metadata and display stats based on the global tiff_selector
        # after you change any values in the selector cell above, run this cell  that to trigger output shapefiles.
        import pathlib
        import marimo as mo
        import rioxarray

        # Stop execution if tiff_selector is not available or has no value
        try:
            mo.stop(not tiff_selector.value, mo.md("No TIFF file selected."))
            _tiff_name = tiff_selector.value
        except NameError:
            # Fallback if the UI element isn't found in current scope
            _tiff_dir = pathlib.Path("tiffs")
            _tiff_files = sorted([f.name for f in _tiff_dir.glob("*.tif*")])
            mo.stop(not _tiff_files, mo.md("No TIFF files found in the directory."))
            _tiff_name = _tiff_files[0]

        # Load the metadata and a low-resolution overview of the selected map
        # We use masked=True to handle nodata values correctly
        _tiff_dir = pathlib.Path("tiffs")
        _path_to_file = _tiff_dir / _tiff_name
        _raster = rioxarray.open_rasterio(_path_to_file, masked=True)

        # Extract basic metadata
        _crs = _raster.rio.crs
        _bounds = _raster.rio.bounds()
        _res = _raster.rio.resolution()

        return mo.vstack([
            mo.md(f"### Metadata for `{_tiff_name}`"),
            mo.hstack([
                mo.stat(label="CRS", value=str(_crs)),
                mo.stat(label="Resolution", value=f"{_res[0]:.6f}, {_res[1]:.6f}"),
                mo.stat(label="Bounds", value=f"L: {_bounds[0]:.2f}, R: {_bounds[2]:.2f}")
            ])
        ])


    _()
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    # Now let's load functions
    """)
    return


@app.cell
def _(measure, mo, np, text_mask, tiff_selector):
    # i think this runs on remnants of the default image when it first loads
    # and then it gets called again later with the current image

    def _():
        from skimage import morphology
        print(f"Running full-resolution analysis for: {tiff_selector.value}")

        # 1. Clean the mask: Remove very thin "noise" and break thin connections
        # This helps separate buildings from roads that they might be touching
        _mask_np = text_mask.values.astype(np.uint8)
        _cleaned_mask = morphology.binary_opening(_mask_np, morphology.square(2))

        # 2. Run analysis on the cleaned mask
        _labels = measure.label(_cleaned_mask)
        _props = measure.regionprops(_labels)

        # Re-initialize masks at full resolution
        text_only_mask = np.zeros(text_mask.shape, dtype=np.uint8)
        building_only_mask = np.zeros(text_mask.shape, dtype=np.uint8)
        line_features_mask = np.zeros(text_mask.shape, dtype=np.uint8)

        for prop in _props:
            _area = prop.area
            _solidity = prop.solidity
            _ecc = prop.eccentricity

            # BUILDINGS: High solidity, compact area
            if 60 <= _area <= 2000 and _solidity > 0.8:
                for coord in prop.coords:
                    building_only_mask[coord[0], coord[1]] = 1

            # ROADS/LINES: Very long, low solidity relative to length
            elif _ecc > 0.98 or (_area > 500 and _solidity < 0.3):
                for coord in prop.coords:
                    line_features_mask[coord[0], coord[1]] = 1

            # TEXT: Small chunks that are horizontal-ish
            elif _area < 150:
                _orientation = abs(prop.orientation)
                if _orientation > 1.3 or _orientation < 0.3:
                    for coord in prop.coords:
                        text_only_mask[coord[0], coord[1]] = 1
        return mo.md("Full-resolution analysis complete with morphological cleaning.")


    _()
    return


@app.cell
def _(
    gpd,
    measure,
    min_building_size,
    mo,
    morphology,
    np,
    pathlib,
    pyproj,
    rioxarray,
    sg,
    threshold_slider,
    tiff_dir,
    tiff_selector,
):
    def _():
        mo.stop(not tiff_selector.value)
        print(f"Running extraction results for: {tiff_selector.value}")

        # 1. Load and Preprocess
        path_to_file = tiff_dir / tiff_selector.value
        raster = rioxarray.open_rasterio(path_to_file, masked=True)
        base_name = pathlib.Path(tiff_selector.value).stem

        # Grayscale & Mask
        if raster.rio.count >= 3:
            grayscale = (0.299 * raster.sel(band=1) + 0.587 * raster.sel(band=2) + 0.114 * raster.sel(band=3))
        else:
            grayscale = raster.squeeze()

        # Binary mask of dark features
        binary_mask = (grayscale < threshold_slider.value).values.astype(np.uint8)
        cleaned_mask = morphology.binary_opening(binary_mask, morphology.square(2))

        # 2. Extract Geometry
        labels = measure.label(cleaned_mask)
        props = measure.regionprops(labels)
        transformer = pyproj.Transformer.from_crs(raster.rio.crs, "EPSG:4326", always_xy=True)

        buildings = []
        # Use min_building_size.value as defined in the previous cell
        for p in props:
            # Heuristic: Buildings are solid blocks of a certain size
            if p.area >= min_building_size.value and p.solidity > 0.8:
                patch = cleaned_mask[p.slice]
                if patch.shape[0] < 2 or patch.shape[1] < 2: continue

                contours = measure.find_contours(patch.astype(float), 0.5)
                if contours:
                    c = contours[0]
                    rows = c[:, 0] + p.slice[0].start
                    cols = c[:, 1] + p.slice[1].start
                    lons_map, lats_map = raster.rio.transform() * (cols, rows)
                    wgs_coords = [transformer.transform(x, y) for x, y in zip(lons_map, lats_map)]

                    if len(wgs_coords) >= 3:
                        buildings.append({
                            "geometry": sg.Polygon(wgs_coords),
                            "area_px": p.area,
                            "solidity": p.solidity
                        })

        # 3. Export
        output_dir = pathlib.Path("outputdata")
        output_dir.mkdir(exist_ok=True)
        shp_path = output_dir / f"{base_name}_buildings.shp"

        if buildings:
            gdf = gpd.GeoDataFrame(buildings, crs="EPSG:4326")
            gdf.to_file(shp_path)
            result_msg = f"Successfully exported {len(buildings)} buildings to `{shp_path}`"
        else:
            result_msg = "No buildings found with current settings."
        return mo.md(f"### Extraction Results\n{result_msg}")


    _()
    return


@app.cell
def _(measure, mo, np, text_mask, tiff_selector):
    # masks try 2
    # if this is just a function definition, I don't know why it runs when it's first loaded. Where is it getting its data?
    def _():
        print(f"Running mask classification (try 2) for: {tiff_selector.value}")
        # Updated mask initialization with geometric heuristics
        # 1. Text: Small, horizontal (orientation near 0 or pi), high eccentricity
        # 2. Small Buildings: ~80 to 200 pixels
        # 3. Industrial/Public: > 200 pixels

        # Initialize masks
        text_only_mask = np.zeros(text_mask.shape, dtype=np.uint8)
        small_building_mask = np.zeros(text_mask.shape, dtype=np.uint8)
        large_building_mask = np.zeros(text_mask.shape, dtype=np.uint8)

        _mask_np = text_mask.values.astype(np.uint8)
        _labels = measure.label(_mask_np)
        _props = measure.regionprops(_labels)

        for prop in _props:
            # Get basic geometry
            _area = prop.area
            _ecc = prop.eccentricity
            # Orientation is in radians from -pi/2 to pi/2
            _orientation = abs(prop.orientation) 

            # 1. HEURISTIC FOR TEXT (Cyrillic labels)
            # Text chunks are usually small and oriented horizontally (near pi/2 or 0 depending on axis)
            # In skimage, horizontal is roughly 1.57 (pi/2)
            _is_horizontal = _orientation > 1.4 or _orientation < 0.2
            if _area < 80 and _is_horizontal:
                for coord in prop.coords:
                    text_only_mask[coord[0], coord[1]] = 1

            # 2. HEURISTIC FOR SMALL BUILDINGS (Villages/Residences)
            elif 60 <= _area <= 200:
                for coord in prop.coords:
                    small_building_mask[coord[0], coord[1]] = 1

            # 3. HEURISTIC FOR LARGE BUILDINGS (Industrial/Public)
            elif _area > 200:
                # We check solidity to ensure it's a 'block' and not a long contour line
                if prop.solidity > 0.5:
                    for coord in prop.coords:
                        large_building_mask[coord[0], coord[1]] = 1

        # Combine buildings for backward compatibility with other cells
        building_only_mask = (small_building_mask | large_building_mask).astype(np.uint8)
        return mo.md(
            f"""
            ### Refined Mask Classification
            - **Text Blobs**: Found {np.count_nonzero(text_only_mask)} pixels
            - **Small Buildings (8x10 range)**: Found {np.count_nonzero(small_building_mask)} pixels
            - **Large/Industrial Buildings**: Found {np.count_nonzero(large_building_mask)} pixels
            """
        )


    _()
    return


@app.cell
def _(
    measure,
    mo,
    np,
    pl,
    pyproj,
    rioxarray,
    text_mask,
    tiff_dir,
    tiff_selector,
):
    def _():
        import shapely.geometry as sg
        print(f"Compiling building polygons for: {tiff_selector.value}")

        # Ensure dependencies are available
        mo.stop(not tiff_selector.value, mo.md("No TIFF selected."))

        # 1. COMPILE BUILDINGS INTO POLYGONS
        buildings_polygons = []

        # Re-extract the building_only_mask from the established logic to avoid NameError
        _mask_np = text_mask.values.astype(np.uint8)
        _labels_init = measure.label(_mask_np)
        _props_init = measure.regionprops(_labels_init)

        _building_only_mask_local = np.zeros(text_mask.shape, dtype=np.uint8)
        for _prop in _props_init:
            # Heuristic for buildings: Solid blocks with high solidity and reasonable area
            if 60 <= _prop.area <= 2000 and _prop.solidity > 0.8:
                for _coord in _prop.coords:
                    _building_only_mask_local[_coord[0], _coord[1]] = 1

        # Label the infrastructure mask again to iterate through individual blobs
        _build_labels = measure.label(_building_only_mask_local)
        _build_props = measure.regionprops(_build_labels)

        # Re-open or use existing raster to get CRS and transform
        _path_to_file = tiff_dir / tiff_selector.value
        _raster_local = rioxarray.open_rasterio(_path_to_file, masked=True)

        # Transformer for coordinate conversion
        _transformer = pyproj.Transformer.from_crs(_raster_local.rio.crs, "EPSG:4326", always_xy=True)

        for prop in _build_props:
            # Extract the binary patch for the current property
            _patch = _building_only_mask_local[prop.slice]

            # measure.find_contours requires at least a 2x2 array.
            # If the bounding box of the building is too small, we skip it.
            if _patch.shape[0] < 2 or _patch.shape[1] < 2:
                continue

            # Find contours for this specific object (returns list of [row, col] arrays)
            # We use level=0.5 for binary masks
            _contours = measure.find_contours(_patch, level=0.5)

            if _contours:
                # Take the longest contour (the outer boundary)
                _contour = _contours[0]

                # Adjust contour coordinates back to global raster pixel coordinates
                _rows = _contour[:, 0] + prop.slice[0].start
                _cols = _contour[:, 1] + prop.slice[1].start

                # Transform pixels to Map CRS then to WGS84
                _lons, _lats = _raster_local.rio.transform() * (_cols, _rows)
                _wgs_coords = [_transformer.transform(x, y) for x, y in zip(_lons, _lats)]

                if len(_wgs_coords) >= 3:
                    buildings_polygons.append({
                        "geometry": sg.Polygon(_wgs_coords),
                        "area_px": prop.area,
                        "eccentricity": prop.eccentricity
                    })

        # Convert to a Geo-ready Polars DataFrame (using WKT for geometry for compatibility)
        if buildings_polygons:
            buildings_final = pl.DataFrame([
                {**p, "geometry": p["geometry"].wkt} for p in buildings_polygons
            ])
        else:
            buildings_final = pl.DataFrame(schema={"geometry": pl.Utf8, "area_px": pl.Float64, "eccentricity": pl.Float64})

        return mo.md(f"### Compiled {len(buildings_final)} Building Polygons")


    _()
    return


@app.cell
def _(folium, mo, pyproj, raster, sampled_buildings):
    def _():
        import base64
        from io import BytesIO
        from PIL import Image
        import numpy as np

        # Ensure dependencies from other cells are available
        mo.stop("sampled_buildings" not in locals(), mo.md("Please run the sampling cell first."))

        def _get_raster_chip(lat_wgs, lon_wgs, buffer_px=150):
            """Extracts a pixel chip from the raster around a WGS84 coordinate."""
            # 1. Transform WGS84 back to Raster CRS to find pixel coordinates
            _back_transformer = pyproj.Transformer.from_crs("EPSG:4326", raster.rio.crs, always_xy=True)
            _target_x, _target_y = _back_transformer.transform(lon_wgs, lat_wgs)

            # 2. Get pixel indices using the raster's inverse transform
            _inv_transform = ~raster.rio.transform()
            _px_col, _px_row = [int(_v) for _v in _inv_transform * (_target_x, _target_y)]

            # 3. Slice the raster (handle boundaries)
            _y_slice = slice(max(0, _px_row - buffer_px), min(raster.rio.height, _px_row + buffer_px))
            _x_slice = slice(max(0, _px_col - buffer_px), min(raster.rio.width, _px_col + buffer_px))

            _chip = raster.isel(y=_y_slice, x=_x_slice)

            # 4. Convert to RGB PIL Image
            if _chip.rio.count == 3:
                _data = _chip.transpose("y", "x", "band").values
                if _data.max() > 1: _data = (_data).astype(np.uint8)
            else:
                _data = _chip.squeeze().values
                # Simple grayscale to RGB for PIL
                _data = np.stack([_data]*3, axis=-1).astype(np.uint8)

            _img = Image.fromarray(_data)

            # 5. Encode to Base64 for Folium Popup
            _buffered = BytesIO()
            _img.save(_buffered, format="PNG")
            _img_str = base64.b64encode(_buffered.getvalue()).decode()
            return f'<img src="data:image/png;base64,{_img_str}" width="200px">'

        # Create the map
        _m_chips = folium.Map(
            location=[sampled_buildings["lat"].mean(), sampled_buildings["lon"].mean()], 
            zoom_start=17, 
            tiles="OpenStreetMap"
        )

        for _row in sampled_buildings.to_dicts():
            # Get the image HTML for the popup
            try:
                _html_img = _get_raster_chip(_row["lat"], _row["lon"], buffer_px=150) # 300px total width/height
                _popup_content = folium.Popup(f"<b>Building Candidate</b><br>{_html_img}<br>Area: {_row['area_px']}px", max_width=250)
            except Exception as e:
                _popup_content = f"Error loading chip: {e}"

            # Add the marker
            folium.Marker(
                location=[_row["lat"], _row["lon"]],
                popup=_popup_content,
                icon=folium.Icon(color="red", icon="home")
            ).add_to(_m_chips)
        return _m_chips


    _()
    return


@app.cell
def _(measure, mo, np, text_mask, tiff_selector):
    # try 3 for buildings vs text
    # this cell has a future warning
    def _():
        print(f"Running refined heuristic classification (try 3) for: {tiff_selector.value}")
        # Updated classification with Solidity and Axis Length heuristics
        # 1. Text: Very small, horizontal, medium solidity.
        # 2. Buildings: Medium to Large, HIGH solidity (rectangular/square).
        # 3. Roads/Lines: Very high eccentricity, LOW solidity, long major axis.

        _mask_np = text_mask.values.astype(np.uint8)
        _labels = measure.label(_mask_np)
        _props = measure.regionprops(_labels)

        # Re-initialize masks
        text_only_mask = np.zeros(text_mask.shape, dtype=np.uint8)
        building_only_mask = np.zeros(text_mask.shape, dtype=np.uint8)
        line_features_mask = np.zeros(text_mask.shape, dtype=np.uint8)

        for prop in _props:
            _area = prop.area
            _ecc = prop.eccentricity
            _solidity = prop.solidity
            _orientation = abs(prop.orientation)
            # axis_major_length is available in newer skimage, otherwise use major_axis_length
            _length = getattr(prop, "axis_major_length", prop.major_axis_length)

            # 1. IDENTIFY ROADS / CONTOURS / BORDERS (Linear features)
            # These are very long and not solid (thin lines)
            if _ecc > 0.97 or (_ecc > 0.93 and _solidity < 0.4):
                for coord in prop.coords:
                    line_features_mask[coord[0], coord[1]] = 1

            # 2. IDENTIFY BUILDINGS (Solid blocks)
            # Buildings have high solidity because they are usually filled rectangles/squares
            elif _area > 60 and _solidity > 0.7:
                for coord in prop.coords:
                    building_only_mask[coord[0], coord[1]] = 1

            # 3. IDENTIFY TEXT (Small, horizontal, non-linear)
            elif _area < 100:
                # Check if it's roughly horizontal (Cyrillic labels)
                _is_horizontal = _orientation > 1.4 or _orientation < 0.2
                if _is_horizontal and _solidity > 0.3:
                    for coord in prop.coords:
                        text_only_mask[coord[0], coord[1]] = 1
        return mo.md(
            f"""
            ### Refined Heuristic Classification
            - **Buildings**: {np.count_nonzero(building_only_mask)} pixels (Filtered by high Solidity)
            - **Text**: {np.count_nonzero(text_only_mask)} pixels (Filtered by Area & Orientation)
            - **Roads/Lines**: {np.count_nonzero(line_features_mask)} pixels (Filtered by Eccentricity)
            """
        )


    _()
    return


@app.cell
def _(folium, gdf, mo):
    def _():


        def _create_map():
            # Using buildings_final (or checking if it exists) from previous processing steps
            mo.stop('buildings_polygons' not in locals() and 'gdf' not in locals(), mo.md("No buildings extracted to visualize."))

            # Use the GeoDataFrame defined in the extraction block
            _gdf = gdf

            # Calculate map center from the GeoDataFrame
            _center_lat = _gdf.geometry.centroid.y.mean()
            _center_lon = _gdf.geometry.centroid.x.mean()

            _m = folium.Map(location=[_center_lat, _center_lon], zoom_start=15, tiles="OpenStreetMap")

            # Add the GeoDataFrame to the map
            folium.GeoJson(
                _gdf,
                name="Extracted Buildings",
                tooltip=folium.GeoJsonTooltip(fields=["area_px", "solidity"], aliases=["Area (px)", "Solidity"]),
                style_function=lambda x: {
                    "fillColor": "red",
                    "color": "darkred",
                    "weight": 1,
                    "fillOpacity": 0.6,
                }
            ).add_to(_m)

            return _m
        return _create_map()


    _()
    return


@app.cell
def _():
    return


@app.cell
def _(raster, tiff_selector):
    def _():
        # is this a cell that will print out the current map no matter where we're at?
        # it is definitely working with the current loaded map on 2/25
        import matplotlib.pyplot as plt
        import numpy as np

        # Ensure the raster is loaded. If 'raster' isn't in locals, 
        # we attempt to load it using the current tiff_selector value.
        try:
            _current_raster = raster
        except NameError:
            import rioxarray
            import pathlib
            _tiff_path = pathlib.Path("tiffs") / tiff_selector.value
            _current_raster = rioxarray.open_rasterio(_tiff_path, masked=True)

        # Determine an appropriate downsample factor to keep the preview size manageable for Marimo.
        # To ensure we stay well below 10MB in the notebook frontend, 
        # we aim for a maximum dimension of around 1200 pixels.
        _max_dim = max(_current_raster.rio.width, _current_raster.rio.height)
        _downsample_factor = max(1, int(_max_dim / 1200))
    
        # Safety floor
        _downsample_factor = max(_downsample_factor, 1)

        _preview_data = _current_raster.isel(
            x=slice(0, -1, _downsample_factor), 
            y=slice(0, -1, _downsample_factor)
        )

        # Prepare the data for plotting
        # If it's 3-band RGB (C, Y, X), we need to transpose it to (Y, X, C) for matplotlib
        if _preview_data.rio.count == 3:
            _plot_data = _preview_data.transpose("y", "x", "band").values
            # Normalize to 0-1 if values are 0-255 (standard for scans)
            if _plot_data.max() > 1:
                _plot_data = _plot_data / 255.0
        else:
            _plot_data = _preview_data.squeeze().values

        plt.figure(figsize=(10, 8))
        plt.imshow(_plot_data, cmap="gray" if _preview_data.rio.count == 1 else None)
        plt.title(f"Visual Preview: {tiff_selector.value} (Downsampled {_downsample_factor}x)")
        plt.axis("off")

        # In marimo, returning the figure object is the standard for display
        return plt.gca()


    _()
    return


@app.cell
def _(measure, np, raster, text_mask, threshold_slider, tiff_selector):
    # This is an analytics cell that runs for a while

    import matplotlib.pyplot as plt

    print(f"Running visual analytics for: {tiff_selector.value}")

    if 'text_mask' not in locals() and 'raster' in locals():
        # Convert to grayscale if it's RGB
        if raster.rio.count == 3:
            _grayscale_local = (
                0.299 * raster.sel(band=1) + 
                0.587 * raster.sel(band=2) + 
                0.114 * raster.sel(band=3)
            )
        else:
            _grayscale_local = raster.squeeze()

        # Use the threshold from the slider or a default value
        _thresh = threshold_slider.value if 'threshold_slider' in locals() else 90
        _text_mask_local = _grayscale_local < _thresh
    else:
        _text_mask_local = text_mask

    _mask_np_display = _text_mask_local.values.astype(np.uint8)
    _labels_display = measure.label(_mask_np_display)
    _props_display = measure.regionprops(_labels_display)

    _text_only_mask_display = np.zeros(_text_mask_local.shape, dtype=np.uint8)
    _building_only_mask_display = np.zeros(_text_mask_local.shape, dtype=np.uint8)

    for _prop in _props_display:
        _area = _prop.area
        _ecc = _prop.eccentricity
        _solidity = _prop.solidity
        _orientation = abs(_prop.orientation)

        # Re-apply the logic used in previous successful analysis to define the masks globally for the plot
        if 60 <= _area <= 2000 and _solidity > 0.8:
            for _coord in _prop.coords:
                _building_only_mask_display[_coord[0], _coord[1]] = 1
        elif _area < 150:
            if _orientation > 1.3 or _orientation < 0.3:
                for _coord in _prop.coords:
                    _text_only_mask_display[_coord[0], _coord[1]] = 1

    _fig_sep, (_ax1, _ax2) = plt.subplots(1, 2, figsize=(16, 8))

    # Downsample for preview
    _ds = 5
    _text_preview = _text_only_mask_display[::_ds, ::_ds]
    _build_preview = _building_only_mask_display[::_ds, ::_ds]

    _ax1.imshow(_text_preview, cmap="Greys")
    _ax1.set_title("Potential Text (Small Blobs)")
    _ax1.axis("off")

    _ax2.imshow(_build_preview, cmap="Greys")
    _ax2.set_title("Infrastructure (Large/Rectilinear)")
    _ax2.axis("off")

    plt.gca()
    return (plt,)


@app.cell
def _(mo, raster, rioxarray, threshold_slider, tiff_dir, tiff_selector):
    # Convert to grayscale if it's RGB, otherwise use the single band
    _raster_source = raster if 'raster' in locals() else rioxarray.open_rasterio(tiff_dir / tiff_selector.value, masked=True)

    if _raster_source.rio.count >= 3:
        grayscale = (
            0.299 * _raster_source.sel(band=1) + 
            0.587 * _raster_source.sel(band=2) + 
            0.114 * _raster_source.sel(band=3)
        )
    else:
        grayscale = _raster_source.squeeze()

    # Create the binary mask where 'True' (1) represents the dark features
    # We invert it so that the features (text/lines) are 1 and background is 0
    text_mask = grayscale < threshold_slider.value

    mo.md(f"Mask created using threshold: {threshold_slider.value}")
    return grayscale, text_mask


@app.cell
def _(buildings_df, mo, pl, raster, rioxarray, tiff_dir, tiff_selector):
    def _():
        import pyproj

        # Ensure raster and buildings_df are available
        try:
            _raster_ref = raster
            _buildings_ref = buildings_df
        except NameError:
            # Fallback to reload if variables aren't in scope
            _path_to_file = tiff_dir / tiff_selector.value
            _raster_ref = rioxarray.open_rasterio(_path_to_file, masked=True)
            # We need the buildings_df produced by the previous extraction steps
            mo.stop("buildings_df" not in locals(), mo.md("Please run the building extraction cell first."))
            _buildings_ref = buildings_df

        # Get the source CRS from the raster metadata
        _src_crs = _raster_ref.rio.crs

        # Create a transformer to WGS84 (Lat/Lon)
        _transformer = pyproj.Transformer.from_crs(_src_crs, "EPSG:4326", always_xy=True)

        # Update buildings_df with WGS84 coordinates
        _wgs84_data = []
        for _row in _buildings_ref.to_dicts():
            # Convert from map CRS to WGS84
            _lon_wgs, _lat_wgs = _transformer.transform(_row["lon"], _row["lat"])
            _wgs84_data.append({
                "area_px": _row["area_px"],
                "eccentricity": _row["eccentricity"],
                "lon": _lon_wgs,
                "lat": _lat_wgs
            })

        # Overwrite buildings_df with the corrected coordinates
        buildings_df_wgs84 = pl.DataFrame(_wgs84_data)

        # Sample 10 buildings
        sampled_buildings = buildings_df_wgs84.sample(n=min(10, len(buildings_df_wgs84)), seed=42)
        return mo.md(f"### Visualizing {len(sampled_buildings)} Sampled Buildings (Reprojected to WGS84)")


    _()
    return


@app.cell
def _(measure, np, plt, text_mask):
    # classify and plot
    _fig_refined, (_ax_t, _ax_s, _ax_l) = plt.subplots(1, 3, figsize=(18, 6))
    _ds = 2  # downsample for preview

    # Re-extract masks from text_mask values for this visualization cell
    # to ensure variables are available globally/outside the previous function scope.
    _mask_np = text_mask.values.astype(np.uint8)
    _labels = measure.label(_mask_np)
    _props = measure.regionprops(_labels)

    _text_only_mask = np.zeros(text_mask.shape, dtype=np.uint8)
    _building_only_mask = np.zeros(text_mask.shape, dtype=np.uint8)

    for _prop in _props:
        _area = _prop.area
        _ecc = _prop.eccentricity
        _solidity = _prop.solidity
        _orientation = abs(_prop.orientation)

        # Simple heuristic to recreate the classification for display
        if _area > 60 and _solidity > 0.7:
            for _coord in _prop.coords:
                _building_only_mask[_coord[0], _coord[1]] = 1
        elif _area < 100:
            _is_horizontal = _orientation > 1.4 or _orientation < 0.2
            if _is_horizontal and _solidity > 0.3:
                for _coord in _prop.coords:
                    _text_only_mask[_coord[0], _coord[1]] = 1

    _ax_t.imshow(_text_only_mask[::_ds, ::_ds], cmap="Greys")
    _ax_t.set_title("Classified: Text")
    _ax_t.axis("off")

    _ax_s.imshow(_building_only_mask[::_ds, ::_ds], cmap="Greys")
    _ax_s.set_title("Classified: Buildings")
    _ax_s.axis("off")

    _ax_l.imshow(text_mask.values[::_ds, ::_ds], cmap="Greys")
    _ax_l.set_title("Original Binary Mask")
    _ax_l.axis("off")

    plt.gca()
    return


@app.cell
def _(measure, np, plt, text_mask):
    # output
    _fig_sep, (_ax1, _ax2, _ax3) = plt.subplots(1, 3, figsize=(18, 6))
    _ds = 4

    # To fix the NameError, we ensure the masks are extracted from the text_mask if they aren't globally available
    _mask_np = text_mask.values.astype(np.uint8)
    _labels = measure.label(_mask_np)
    _props = measure.regionprops(_labels)

    _building_mask = np.zeros(text_mask.shape, dtype=np.uint8)
    _text_mask = np.zeros(text_mask.shape, dtype=np.uint8)
    _line_mask = np.zeros(text_mask.shape, dtype=np.uint8)

    for _prop in _props:
        _area = _prop.area
        _ecc = _prop.eccentricity
        _solidity = _prop.solidity
        _orientation = abs(_prop.orientation)

        if _ecc > 0.97 or (_ecc > 0.93 and _solidity < 0.4):
            for _coord in _prop.coords:
                _line_mask[_coord[0], _coord[1]] = 1
        elif _area > 60 and _solidity > 0.7:
            for _coord in _prop.coords:
                _building_mask[_coord[0], _coord[1]] = 1
        elif _area < 100:
            _is_horizontal = _orientation > 1.4 or _orientation < 0.2
            if _is_horizontal and _solidity > 0.3:
                for _coord in _prop.coords:
                    _text_mask[_coord[0], _coord[1]] = 1

    _ax1.imshow(_building_mask[::_ds, ::_ds], cmap="Greys")
    _ax1.set_title("Identified Buildings (Solid Blocks)")
    _ax1.axis("off")

    _ax2.imshow(_text_mask[::_ds, ::_ds], cmap="Greys")
    _ax2.set_title("Potential Text (Small/Horizontal)")
    _ax2.axis("off")

    _ax3.imshow(_line_mask[::_ds, ::_ds], cmap="Greys")
    _ax3.set_title("Roads & Contours (High Eccentricity)")
    _ax3.axis("off")

    plt.gca()
    return


@app.cell
def slowocrcell(grayscale, measure, mo, np, pl, text_mask, tiff_selector):
    # this one takes a bit
    # this cell purports do to english and cyrillic OCR
    # when running in a fresh kernel, it gets called for the first time by one of the functions
    # it also gets called by an export shapefile cell. that doesn't seem right
    import pytesseract
    from PIL import Image

    # Print the name of the current map
    print(f"Running OCR for: {tiff_selector.value}")

    # Re-extract the text_only_mask from the grayscale data if not globally available
    # based on the classification logic from previous cells
    _mask_np = text_mask.values.astype(np.uint8)
    _labels = measure.label(_mask_np)
    _props = measure.regionprops(_labels)

    _text_only_mask_local = np.zeros(text_mask.shape, dtype=np.uint8)
    for _prop in _props:
        _area = _prop.area
        _orientation = abs(_prop.orientation)
        _solidity = _prop.solidity

        # Heuristic for text: Small chunks, horizontal-ish
        if _area < 150:
            _is_horizontal = _orientation > 1.3 or _orientation < 0.2
            if _is_horizontal and _solidity > 0.3:
                for _coord in _prop.coords:
                    _text_only_mask_local[_coord[0], _coord[1]] = 1

    # We use the original grayscale image but mask it with our text-only mask
    # to remove background noise and lines
    _text_only_image = grayscale.values * _text_only_mask_local

    # Convert to a PIL Image for Tesseract
    _pil_image = Image.fromarray((_text_only_image * 255).astype(np.uint8))

    # Perform OCR (using Russian and English languages)
    # --psm 11 looks for sparse text
    ocr_data = pytesseract.image_to_data(_pil_image, lang='rus+eng', output_type=pytesseract.Output.DATAFRAME)

    # Convert to polars using the dictionary approach to avoid pyarrow dependency issues with mixed types
    found_text = pl.from_dicts(ocr_data.to_dict('records')).filter(pl.col("conf") > 50)
    mo.ui.table(found_text)
    return (found_text,)


@app.cell
def _(
    measure,
    mo,
    np,
    pl,
    raster,
    rioxarray,
    text_mask,
    tiff_dir,
    tiff_selector,
):
    # Extract properties of the "Infrastructure" blobs
    # Re-extract the building_only_mask from the logic established in previous cells
    _mask_np = text_mask.values.astype(np.uint8)
    _labels_init = measure.label(_mask_np)
    _props_init = measure.regionprops(_labels_init)

    # Ensure raster is loaded if not already in scope
    try:
        _raster_ref = raster
    except NameError:
        _path_to_file = tiff_dir / tiff_selector.value
        _raster_ref = rioxarray.open_rasterio(_path_to_file, masked=True)

    _building_only_mask_local = np.zeros(text_mask.shape, dtype=np.uint8)
    for _prop in _props_init:
        # Heuristic for buildings: Solid blocks with high solidity and reasonable area
        if 60 <= _prop.area <= 2000 and _prop.solidity > 0.8:
            for _coord in _prop.coords:
                _building_only_mask_local[_coord[0], _coord[1]] = 1

    _build_labels = measure.label(_building_only_mask_local)
    _build_props = measure.regionprops(_build_labels)

    buildings_data = []
    for proposed in _build_props:
        # Get pixel centroid
        row, col = proposed.centroid

        # Transform pixel coordinates to CRS coordinates (Lat/Lon or Meters)
        lon, lat = _raster_ref.rio.transform() * (col, row)

        buildings_data.append({
            "area_px": proposed.area,
            "eccentricity": proposed.eccentricity,
            "lon": lon,
            "lat": lat
        })

    buildings_df = pl.DataFrame(buildings_data)
    mo.ui.table(buildings_df)
    return (buildings_df,)


@app.cell
def _():
    return


@app.cell
def _():
    return


@app.cell
def _(
    found_text,
    measure,
    mo,
    np,
    pathlib,
    pyproj,
    rioxarray,
    text_mask,
    tiff_dir,
    tiff_selector,
):
    def _():
        # save buildings, text, and roads shapefiles named after the source TIFF
        import geopandas as gpd
        import os
        import shapely.geometry as sg
        print(f"Exporting data for: {tiff_selector.value}")

        # Ensure necessary objects are available
        mo.stop(not tiff_selector.value, mo.md("No TIFF selected for export."))

        _base_name = pathlib.Path(tiff_selector.value).stem

        # Re-load raster to avoid NameError if not globally available
        _path_to_file = tiff_dir / tiff_selector.value
        _raster_export = rioxarray.open_rasterio(_path_to_file, masked=True)

        # Create an output directory if it doesn't exist
        output_dir = "outputdata"
        os.makedirs(output_dir, exist_ok=True)

        # 1. Compile Features
        _mask_np_export = text_mask.values.astype(np.uint8)
        _labels_init_export = measure.label(_mask_np_export)
        _props_init_export = measure.regionprops(_labels_init_export)

        _buildings_polygons = []
        _roads_polygons = []

        _transformer_export = pyproj.Transformer.from_crs(_raster_export.rio.crs, "EPSG:4326", always_xy=True)

        for _prop in _props_init_export:
            _is_building = 60 <= _prop.area <= 2000 and _prop.solidity > 0.8
            _is_road = _prop.eccentricity > 0.97 or (_prop.area > 500 and _prop.solidity < 0.3)

            if not (_is_building or _is_road):
                continue

            _patch = _mask_np_export[_prop.slice]
            if _patch.shape[0] < 2 or _patch.shape[1] < 2:
                continue

            _contours = measure.find_contours(_patch.astype(float), level=0.5)
            if _contours:
                _contour = _contours[0]
                _rows = _contour[:, 0] + _prop.slice[0].start
                _cols = _contour[:, 1] + _prop.slice[1].start
                _lons, _lats = _raster_export.rio.transform() * (_cols, _rows)
                _wgs_coords = [_transformer_export.transform(x, y) for x, y in zip(_lons, _lats)]

                if len(_wgs_coords) >= 3:
                    _geom = sg.Polygon(_wgs_coords)
                    _entry = {"geometry": _geom, "area_px": _prop.area, "solidity": _prop.solidity}
                    if _is_building:
                        _buildings_polygons.append(_entry)
                    else:
                        _roads_polygons.append(_entry)

        # Export Buildings
        _buildings_path = os.path.join(output_dir, f"{_base_name}_buildings.shp")
        if _buildings_polygons:
            gpd.GeoDataFrame(_buildings_polygons, crs="EPSG:4326").to_file(_buildings_path)
            _build_stat = mo.stat(label="Buildings", value=f"{len(_buildings_polygons)} features")
        else:
            _build_stat = mo.stat(label="Buildings", value="None found")

        # Export Roads
        _roads_path = os.path.join(output_dir, f"{_base_name}_roads.shp")
        if _roads_polygons:
            gpd.GeoDataFrame(_roads_polygons, crs="EPSG:4326").to_file(_roads_path)
            _road_stat = mo.stat(label="Roads/Lines", value=f"{len(_roads_polygons)} features")
        else:
            _road_stat = mo.stat(label="Roads/Lines", value="None found")

        # 2. Export Text (Labels)
        _text_path = os.path.join(output_dir, f"{_base_name}_labels.shp")
        if 'found_text' in locals() or 'found_text' in globals():
            _text_src = found_text if 'found_text' in locals() else found_text
            _gdf_text = gpd.GeoDataFrame(
                _text_src.to_dicts(),
                geometry=gpd.points_from_xy(_text_src["left"], _text_src["top"]),
                crs="EPSG:4326"
            )
            _gdf_text.to_file(_text_path)
            _text_stat = mo.stat(label="Text Labels", value=f"{len(_text_src)} labels")
        else:
            _text_stat = mo.stat(label="Text Labels", value="Run OCR first")

        return mo.vstack([
            mo.md(f"### Data Exported Successfully for `{_base_name}`"),
            mo.hstack([
                _build_stat,
                _road_stat,
                _text_stat
            ]),
            mo.md(f"Files saved in `{output_dir}/` as `{_base_name}_*.shp`.")
        ])


    _()
    return


@app.cell
def _():
    return


@app.cell
def _(mo, raster, rioxarray, tiff_dir, tiff_selector):
    def _():
        # Convert to grayscale if it's RGB, otherwise use the single band
        _raster_source = raster if 'raster' in locals() else rioxarray.open_rasterio(tiff_dir / tiff_selector.value, masked=True)

        if _raster_source.rio.count == 3:
            # Standard grayscale conversion: 0.299R + 0.587G + 0.114B
            grayscale = (
                0.299 * _raster_source.sel(band=1) + 
                0.587 * _raster_source.sel(band=2) + 
                0.114 * _raster_source.sel(band=3)
            )
        else:
            grayscale = _raster_source.squeeze()

        # Thresholding: In Soviet maps, text/buildings are usually dark (low values)
        # Adjust the threshold (e.g., 100) based on the scan quality
        threshold_slider = mo.ui.slider(0, 255, value=90, label="Darkness Threshold (Lower = Darker)")
        return threshold_slider


    _()
    return


@app.cell
def _(mo):
    def _():
        # building size selector from the original app. don't run me
        from skimage import measure, morphology
        import polars as pl

        # Convert the boolean mask to an integer label image
        # Objects with 150+ pixels are likely buildings/infrastructure
        # smallest structures on the map are about 40 pixels
        min_building_size = mo.ui.number(start=10, stop=1000, value=400, label="Min Building Size (pixels)")
        return min_building_size


    _()
    return


if __name__ == "__main__":
    app.run()