Update app.py

app.py

@@ -3,7 +3,10 @@ import torch
 from diffusers import DiffusionPipeline
 from datasets import load_dataset
 from PIL import Image
-
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import numpy as np
+import os
 
 # Configurações do dispositivo para uso apenas da CPU
 device = "cpu"
@@ -14,7 +17,7 @@ pipe = DiffusionPipeline.from_pretrained(model_repo_id)
 pipe = pipe.to(device)
 
 # Carregando o dataset do Hugging Face
-dataset = load_dataset("LEIDIA/Data_Womleimg")
+dataset = load_dataset("LEIDIA/Data_Womleimg", split="train")
 
 
 # Parâmetros para carregar o dataset personalizado
@@ -39,6 +42,40 @@ dataset_descriptions = [
     "A girl wearing a form-fitting black leather top, with long pink hair cascading down, creating a striking contrast in a neutral background."
 ]
 
+
+# Criar transformações usando Albumentations
+albumentations_transform = A.Compose([
+    A.Resize(512, 512),  # Redimensionar para o tamanho desejado
+    A.HorizontalFlip(p=0.5),  # Flip horizontal com probabilidade de 50%
+    A.RandomBrightnessContrast(p=0.2),  # Ajuste de brilho e contraste
+    A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),  # Adicionar ruído gaussiano
+])
+
+# Caminho temporário para salvar imagens transformadas
+transformed_images_dir = "./transformed_images"
+os.makedirs(transformed_images_dir, exist_ok=True)
+
+# Aplicar transformações no dataset
+def apply_transformations(example, idx):
+    # Carregar imagem usando PIL
+    image = Image.open(example["image"]).convert("RGB")
+    image_np = np.array(image)
+
+    # Aplicar Albumentations
+    transformed = albumentations_transform(image=image_np)
+    transformed_image = transformed["image"]
+
+    # Salvar imagem transformada
+    transformed_image_pil = Image.fromarray(transformed_image)
+    transformed_path = os.path.join(transformed_images_dir, f"transformed_{idx}.png")
+    transformed_image_pil.save(transformed_path)
+
+    return {"image": transformed_path}
+
+# Mapear transformações no dataset
+dataset = dataset.map(apply_transformations, with_indices=True)
+
+
 # Definir parâmetros padrão para geração rápida
 DEFAULT_PROMPT = "A beautiful brunette woman wearing a blue leather pants "
 DEFAULT_INFERENCE_STEPS = 10