generate_hazards.py

from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import os

os.makedirs("synthetic/generated_hazards", exist_ok=True)

def generate_unsafe_image(prompt="construction worker without hard hat, near crane, dangerous zone, realistic", num_images=1):
    pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
    pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

    images = pipe(prompt, num_images_per_prompt=num_images).images

    for i, img in enumerate(images):
        path = f"synthetic/generated_hazards/hazard_{i}.png"
        img.save(path)
        print(f"Generated: {path}")

    return images[0]