Update app.py
Browse files
app.py
CHANGED
@@ -3,8 +3,6 @@ import torch
|
|
3 |
from diffusers import DiffusionPipeline
|
4 |
from datasets import load_dataset
|
5 |
from PIL import Image
|
6 |
-
import albumentations as A
|
7 |
-
from albumentations.pytorch import ToTensorV2
|
8 |
import numpy as np
|
9 |
import os
|
10 |
|
@@ -59,42 +57,6 @@ dataset_descriptions = [
|
|
59 |
]
|
60 |
|
61 |
|
62 |
-
# Criar transformações usando Albumentations
|
63 |
-
albumentations_transform = A.Compose([
|
64 |
-
A.Resize(512, 512), # Redimensionar para o tamanho desejado
|
65 |
-
A.HorizontalFlip(p=0.5), # Flip horizontal com probabilidade de 50%
|
66 |
-
A.RandomBrightnessContrast(p=0.2), # Ajuste de brilho e contraste
|
67 |
-
A.GaussNoise(var_limit=(10.0, 50.0), p=0.3), # Adicionar ruído gaussiano
|
68 |
-
])
|
69 |
-
|
70 |
-
# Caminho temporário para salvar imagens transformadas
|
71 |
-
transformed_images_dir = "./transformed_images"
|
72 |
-
os.makedirs(transformed_images_dir, exist_ok=True)
|
73 |
-
|
74 |
-
# Aplicar transformações no dataset
|
75 |
-
def apply_transformations(example, idx):
|
76 |
-
# Verificar se a imagem é um objeto PIL ou caminho
|
77 |
-
if isinstance(example["image"], Image.Image):
|
78 |
-
image = example["image"] # Imagem já está carregada
|
79 |
-
else:
|
80 |
-
image = Image.open(example["image"]).convert("RGB") # Caso seja um caminho
|
81 |
-
|
82 |
-
# Converter imagem para numpy array
|
83 |
-
image_np = np.array(image)
|
84 |
-
|
85 |
-
# Aplicar Albumentations
|
86 |
-
transformed = albumentations_transform(image=image_np)
|
87 |
-
transformed_image = transformed["image"]
|
88 |
-
|
89 |
-
# Converter de volta para PIL
|
90 |
-
transformed_image_pil = Image.fromarray(transformed_image)
|
91 |
-
|
92 |
-
# Salvar imagem transformada em uma pasta local
|
93 |
-
transformed_path = os.path.join(transformed_images_dir, f"transformed_{idx}.png")
|
94 |
-
transformed_image_pil.save(transformed_path)
|
95 |
-
|
96 |
-
return {"image": transformed_path}
|
97 |
-
|
98 |
|
99 |
# Mapear transformações no dataset
|
100 |
dataset = dataset.map(apply_transformations, with_indices=True)
|
|
|
3 |
from diffusers import DiffusionPipeline
|
4 |
from datasets import load_dataset
|
5 |
from PIL import Image
|
|
|
|
|
6 |
import numpy as np
|
7 |
import os
|
8 |
|
|
|
57 |
]
|
58 |
|
59 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
60 |
|
61 |
# Mapear transformações no dataset
|
62 |
dataset = dataset.map(apply_transformations, with_indices=True)
|