Create preprocess.py

preprocess.py

@@ -0,0 +1,38 @@
+import os
+from PIL import Image
+from transformers import AutoProcessor, BlipForConditionalGeneration
+
+def process_dataset(zip_path, output_dir, generate_captions=True):
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Descompacta dataset
+    import zipfile
+    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
+    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)
+    
+    return output_dir