generate_tree_images/generate_tree_images.py

import os
import rasterio
import geopandas as gpd
from shapely.geometry import box
from rasterio.mask import mask
from PIL import Image
from PIL import ImageOps
import numpy as np
import warnings
from rasterio.errors import NodataShadowWarning
import sys

warnings.filterwarnings("ignore", category=NodataShadowWarning)

def cut_trees(output_dir, geojson_path, tif_path):
    # create output directory if it doesnt exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)


    # Load the GeoDataFrame

    gdf = gpd.read_file(geojson_path)


    # Clear the terminal screen
    # os.system('cls' if os.name == 'nt' else 'clear')

    # Open the .tif file
    with rasterio.open(tif_path) as src:
        # Get the bounds of the .tif image
        tif_bounds = box(*src.bounds)
        tif_bounds = gpd.GeoDataFrame(geometry=[tif_bounds], crs=gdf.crs)
        tif_bounds = tif_bounds['geometry'].iloc[0]
        # Get the CRS (Coordinate Reference System) of the .tif image
        #tif_crs = src.crs
        
        # Reproject the GeoDataFrame to the CRS of the .tif file
       # gdf = gdf.to_crs(tif_crs)
        #print(tif_bounds.crs.to_epsg())
        #print(gdf.crs.to_epsg())

        # Loop through each polygon in the GeoDataFrame
        N = len(gdf)
        n = int(N/10)
        print(f"Processing {N} polygons...")
        image_counter = 0
        for idx, row in gdf.iterrows():
            if idx % n == 0:
                progress = f"{round(idx/N*100)} % complete --> {idx}/{N}"
                sys.stdout.write('\r' + progress)
                sys.stdout.flush()

            # Extract the geometry (polygon)
            geom = row['geometry']
            name = row['id']

    
            
            # Check if the polygon intersects the image bounds
            if geom.intersects(tif_bounds):
                # Create a mask for the current polygon
                out_image, out_transform = mask(src, [geom], crop=True)
                
                # Convert the masked image to a numpy array
                out_image = out_image.transpose(1, 2, 0)  # rearrange dimensions for PIL (H, W, C)

                # Ensure the array is not empty
                if out_image.size == 0:
                    print("Empty image")
                    gdf.drop(idx, inplace=True)
                    message = f"{round(idx/N*100)} % complete --> {idx}/{N} | Polygon {idx} resulted in an empty image and will be skipped."
                    sys.stdout.write('\r' + message)
                    sys.stdout.flush()
                    continue
                
                # Remove the zero-padded areas (optional)
                mask_array = (out_image[:, :, 0] != src.nodata)
                non_zero_rows = np.any(mask_array, axis=1)
                non_zero_cols = np.any(mask_array, axis=0)

                # Ensure there are non-zero rows and columns
                if not np.any(non_zero_rows) or not np.any(non_zero_cols):
                    print("Non zero rows or columns")
                    gdf.drop(idx, inplace=True)
                    message = f"{round(idx/N*100)} % complete --> {idx}/{N} | Polygon {idx} resulted in an invalid image area and will be skipped."
                    sys.stdout.write('\r' + message)
                    sys.stdout.flush()
                    continue

                out_image = out_image[non_zero_rows][:, non_zero_cols]

                # Convert to a PIL Image and save as PNG
                out_image = Image.fromarray(out_image.astype(np.uint8))  # Ensure correct type for PIL
                output_path = os.path.join(output_dir, f'tree_{name}.png')
                out_image.save(output_path)
                image_counter += 1
            else:
                gdf.drop(idx, inplace=True)
                print("Does not intersect")
                message = f"{round(idx/N*100)} % complete --> {idx}/{N} | Polygon {idx} is outside the image bounds and will be skipped."
                sys.stdout.write('\r' + message)
                sys.stdout.flush()
    print(len(gdf))
    gdf.to_file(geojson_path, driver='GeoJSON')
    print(f'\n {image_counter}/{N} Tree images have been successfully saved in the "detected_trees" folder.')


def resize_images(input_folder, output_folder, target_size):
    # Create the output folder if it doesn't exist
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    counter = 0
    # Loop through all files in the input folder
    for filename in os.listdir(input_folder):
        if filename.endswith('.png'):  # Check for PNG files
            # Open image
            with Image.open(os.path.join(input_folder, filename)) as img:
                # Resize image while preserving aspect ratio
                #print("Original image size: ", img.size)

                img.thumbnail(target_size, Image.LANCZOS)

                if img.size[0] < target_size[0] or img.size[1] < target_size[1]:
        
                    # Calculate padding dimensions
                    pad_width = target_size[0] - img.size[0]
                    pad_height = target_size[1] - img.size[1]
        
                    # Calculate padding
                    padding = (pad_width // 2, pad_height // 2, pad_width - (pad_width // 2), pad_height - (pad_height // 2))
        
                    # Pad the image
                    img = ImageOps.expand(img, padding, fill=(0, 0, 0))
                
                #print ("Resized image size: ", img.size)


                # Calculate paste position to center image in canvas
                paste_pos = ((target_size[0] - img.size[0]) // 2, (target_size[1] - img.size[1]) // 2)
                #print("Paste position: ", paste_pos)
                # Create a new blank canvas with the target size and black background
                new_img = Image.new("RGBA", target_size, (0, 0, 0, 255))
                img = img.convert("RGBA")
                # Paste resized image onto the canvas
                new_img.paste(img, paste_pos, img)
                # Convert to RGB to remove transparency by merging with black background
                new_img = new_img.convert("RGB")
                # Save resized image to output folder
                new_img.save(os.path.join(output_folder, filename))
                
            counter += 1
            # Display the counter
            if counter % 100 == 0:
                message = f"Processed {counter} images"
                print(message, end='\r')
    
    # Final message after processing all images
    print(f"Processed a total of {counter} images.")


# THIS IS THE FUNCTION TO IMPORT
def generate_tree_images(geojson_path, tif_path, target_size = (224, 224)):
    """
    INPUT: geojson path, tif_path that contain the trees, optional target_size of the resulting images
    
    RETURNS: nothing
    
    Action: It creates two folders: + "detected trees" --> the cut tree images
                                    + "tree_images" --> the processed cut tree images, ready to use for species recognition
    """


    # Set input and output folders
    folder_cut_trees = "detected_trees"
    folder_finished_images = "tree_images"
    # Set target size (width, height)    
    cut_trees(geojson_path = geojson_path, tif_path = tif_path, output_dir = folder_cut_trees)
    resize_images(input_folder = folder_cut_trees, output_folder = folder_finished_images, target_size = target_size)