| from typing import Literal |
|
|
| import os |
| import torch |
| from itertools import batched |
| import click |
| from PIL import Image |
| from tqdm import tqdm |
| import json |
|
|
| from diffusers import StableDiffusionImg2ImgPipeline |
|
|
|
|
| PROMPTS = { |
| "rain": "a photo of a city street in the rain", |
| "snow": "a photo of a city street in the snow", |
| "fog": "a photo of a city street in the fog", |
| "night": "a photo of a city street at night", |
| } |
|
|
|
|
| @click.command() |
| @click.option("--swim_dir", default="datasets/swim_data") |
| @click.option("--output") |
| @click.option("--type") |
| @click.option("--batch_size", default=1) |
| @click.option("--no_night", is_flag=True, default=False) |
| def sdedit( |
| swim_dir: str, |
| output: str, |
| type: Literal["rain", "snow", "fog", "night"], |
| batch_size: int, |
| no_night: bool, |
| ): |
| device = "cuda" |
| model_id_or_path = "runwayml/stable-diffusion-v1-5" |
| pipe = StableDiffusionImg2ImgPipeline.from_pretrained( |
| model_id_or_path, torch_dtype=torch.float16 |
| ) |
| pipe = pipe.to(device) |
|
|
| os.makedirs(output, exist_ok=True) |
|
|
| with open(f"{swim_dir}/val/labels.json") as f: |
| labels = json.load(f) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| clear_images = [ |
| label["name"] |
| for label in labels |
| if label["weather"] == "clear" |
| and (not no_night or label["timeofday"] != "night") |
| ] |
|
|
| for names in tqdm(list(batched(clear_images, batch_size)), desc="Generating"): |
| images = [Image.open(f"{swim_dir}/val/images/{name}") for name in names] |
| prompts = [PROMPTS[type]] * len(images) |
| processed_images = pipe(images, prompts, device=device) |
|
|
| for name, processed_image in zip(names, processed_images): |
| processed_image.save(f"{output}/{name}") |
|
|
|
|
| if __name__ == "__main__": |
| sdedit() |
|
|
