RubenGres / Seg2Sat

Using StableDiffusion and ControlNet to generate synthetic aerial images

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Seg2Sat - Segmentation to aerial view using pretrained diffuser models

Seg2Sat explores the potential of diffusion algorithms such as StableDiffusion and ControlNet to generate aerial images based on terrain segmentation data. The dataset is derived from IGN's FLAIR (French Land cover from Aerospace Imagery), which provides land cover information for various regions in France. The goal is to create photorealistic synthetic aerial images or stylized aerial imagery.

This could have many uses like dataset creation/augmentation or making maps for your next ungeon and dragon session. Outpainting and inpainting are also possible and yield convincing results.

Example photorealistic images

These images are derived from the same segmentation mask in the first column and a fixed prompt: "Aerial view of 101 avenue de l'Espinet, Toulouse, Haute-Garonne, France". The different image are obtained by varying the seed in the random number generator.

Image Segmentation Generated Images
Legend from the FLAIR dataset documentation
Class Value Freq.-train (%) Freq.-test (%)
building 1 8.14 8.6
pervious surface 2 8.25 7.34
impervious surface 3 13.72 14.98
bare soil 4 3.47 4.36
water 5 4.88 5.98
coniferous 6 2.74 2.39
deciduous 7 15.38 13.91
brushwood 8 6.95 6.91
vineyard 9 3.13 3.87
herbaceous vegetation 10 17.84 22.17
agricultural land 11 10.98 6.95
plowed land 12 3.88 2.25
swimming pool 13 0.03 0.04
snow 14 0.15 -
clear cut 15 0.15 0.01
mixed 16 0.05 -
ligneous 17 0.01 0.03
greenhouse 18 0.12 0.2
other 19 0.14 -

Example stylized images

Since this is derived from a general text-to-image model it is also possible to guide the generation further by modifying the prompt. Prompting can be used for guiding image style, seasonality, and architechture.

The following images were generated using the prompt "<style>, aerial view of 31 Rue Molière, France and varying the syle parameter. The examples were cherry picked with varying seed:

Image Segmentation lego brick pencil sketch oil on canvas stained glass

How to use

Drawing2Map web user interface

This project also comes with a webui to draw in the browser and generate images on the fly

drawing

  • The webui can be used through this link: https://rubengr.es/Seg2Sat
  • You can also run it locally: cd ui && pip install -r requirements.txt && python start_webui.py
  • or with Google Colab:

Diffusers library

from diffusers import StableDiffusionControlNetPipeline, StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
from diffusers.utils import load_image
from PIL import Image
import torch

dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
device = "cuda:0" if torch.cuda.is_available() else "cpu"

controlnet = ControlNetModel.from_pretrained("rgres/Seg2Sat-sd-controlnet", torch_dtype=dtype)

pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1-base", controlnet=controlnet, torch_dtype=dtype, safety_checker=None
).to("cuda")

segmentation = Image.open("segmentation.png")

image = pipe(
   prompt="Aerial view of 31 Rue Molière, Paris, France.",
   num_inference_steps=20, image=image
).images[0]

image.show()

Download the model

The trained weights for the ControlNet model are available on HuggingFace here

About the dataset

To facilitate the use of StableDiffusion, a custom dataset has been derived from the FLAIR dataset. The TIFF images have been converted to PNG format for convenience, and the image segmentation process incorporates the recommended Look-Up Table (LUT) from the dataset documentation. Location data and descriptions have been gathered from the TIF files and the OpenStreetMap API to enrich the dataset, enabling the creation of meaningful labels for each image.

Each label is in the form Aerial view of <OSM display_name>. <classes>
e.g. Aerial view of Pays de la Loire, France métropolitaine, 49260, France. brushwood, deciduous, herbaceous vegetation

The complete dataset, consisting of over 61,000 images, can be accessed on HuggingFace here.

File Structure

This github repo consists of the following directories:

  • controlnet_training/: Contains everything needed to train the ControlNet model with validation images.
  • dataset_preparation/: Contains the code used to convert the FLAIR#1 dataset.
  • endpoint/: Code for the hugging face endpoint
  • ui/: Contains the code for the complete interface allowing you to draw from your browser
  • images/: All images used in the README.md file
  • examples/: Examples scripts for using this model with the diffusers library

Limitations

One limitation of this project is that the FLAIR dataset covers France only. As a result, the learned representation and generated images may be biased towards French terrain. The model might struggle to accurately render other types of terrain outside of France. Future improvements could involve expanding the dataset to include a more diverse range of regions and landscapes.

Also theses results are obtained using a ControlNet on the base StableDiffusion 2.1 model. More details could be obtained in the generated images by finetuning a diffusion model on aerial imaging.

About

Using StableDiffusion and ControlNet to generate synthetic aerial images


Languages

Language:Jupyter Notebook 59.1%Language:JavaScript 26.8%Language:Python 6.3%Language:HTML 4.4%Language:Svelte 2.1%Language:Shell 0.8%Language:TypeScript 0.3%Language:Dockerfile 0.1%Language:Makefile 0.1%Language:CSS 0.0%