Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import json
 from diffusers import DiffusionPipeline
 from datasets import load_dataset
 from PIL import Image
@@ -16,16 +17,18 @@ pipe = pipe.to(device)
 # Carregar o dataset do Hugging Face
 dataset = load_dataset("LEIDIA/Data_Womleimg", split="train")  # Carrega o dataset de imagens
 
-# Exibir os primeiros exemplos para entender a estrutura
-print(dataset[0])  # Mostra o primeiro exemplo do dataset
+# Localizar o arquivo JSON dentro do dataset
+json_path = dataset["train"].cache_files[0]["filename"].replace("data.arrow", "Wom_dataset.json")
+
+# Ler as descrições do JSON
+with open(json_path, "r") as f:
+    descriptions = json.load(f)
 
 # Acessando as descrições corretamente (assumindo que estejam no arquivo JSON)
-descriptions = {}
-for idx, item in enumerate(dataset):
-    image_id = idx  # Usar o índice como chave única
-    # Aqui você pode acessar as descrições da maneira correta, dependendo de como os dados estão estruturados
-    description = item.get("Wom_dataset", "Descrição não encontrada")  # Exemplo de campo de descrição
-    descriptions[image_id] = description  # Mapeando o índice da imagem para a descrição
+for idx, item in enumerate(dataset["train"]):
+    image = item["image"]  # A imagem
+    description = descriptions.get(item["id"], "Descrição não encontrada")
+    print(f"ID da imagem: {idx}, Descrição: {description}")
 
 # Exibe algumas amostras com suas descrições
 for i, (image_id, description) in enumerate(descriptions.items()):