Update preprocess.py

preprocess.py

@@ -1,6 +1,7 @@
+# preprocess.py
 import os
 from PIL import Image
-from transformers import AutoProcessor, BlipForConditionalGeneration
+from transformers import BlipProcessor, BlipForConditionalGeneration
 
 def process_dataset(zip_path, output_dir, generate_captions=True):
     os.makedirs(output_dir, exist_ok=True)
@@ -10,29 +11,27 @@ def process_dataset(zip_path, output_dir, generate_captions=True):
     with zipfile.ZipFile(zip_path, 'r') as zip_ref:
         zip_ref.extractall(output_dir)
     
-    # Gera captions com BLIP multilíngue
-    if generate_captions:
-        processor = AutoProcessor.from_pretrained("microsoft/blip-image-captioning-base")
-        model = BlipForConditionalGeneration.from_pretrained("microsoft/blip-image-captioning-base")
-        
-        for img_name in os.listdir(output_dir):
-            if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
-                img_path = os.path.join(output_dir, img_name)
-                image = Image.open(img_path).convert('RGB')
-                
-                inputs = processor(images=image, return_tensors="pt")
-                outputs = model.generate(**inputs, max_new_tokens=50)
-                caption = processor.decode(outputs[0], skip_special_tokens=True)
-                
-                with open(img_path.replace('.jpg', '.txt').replace('.png', '.txt'), 'w') as f:
-                    f.write(caption)
-    
-    # Redimensiona imagens
+    # Carrega BLIP (em inglês — modelo oficial da Salesforce)
+    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+
+    # Processa imagens
     for img_name in os.listdir(output_dir):
         if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
             img_path = os.path.join(output_dir, img_name)
             image = Image.open(img_path).convert('RGB')
-            image = image.resize((512, 512), Image.LANCZOS)
-            image.save(img_path)
+            
+            # Redimensiona para evitar erros de memória
+            image.thumbnail((512, 512), Image.LANCZOS)
+            image.save(img_path)  # Salva imagem redimensionada
+            
+            if generate_captions:
+                inputs = processor(image, return_tensors="pt")
+                outputs = model.generate(**inputs, max_new_tokens=50)
+                caption = processor.decode(outputs[0], skip_special_tokens=True)
+                
+                txt_path = os.path.splitext(img_path)[0] + ".txt"
+                with open(txt_path, "w", encoding="utf-8") as f:
+                    f.write(caption)
     
     return output_dir