File size: 3,233 Bytes
8a00d0d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import os
import shutil
from PIL import Image
from transformers.image_utils import load_image
import sys
sys.path.append('..')
from vision_encoder import ideficsV3
from tqdm import tqdm
class VisionPreprocessor:
def __init__(self, device=None, param_dtype=torch.float32):
self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
self.param_dtype = param_dtype
# Initialize and freeze the vision encoder
self.vision_encoder = ideficsV3("HuggingFaceTB/SmolVLM-Instruct").eval().to(self.device)
for param in self.vision_encoder.parameters():
param.requires_grad = False
def load_image(self, image_path):
"""Load an image using PIL without preprocessing."""
image = load_image(image_path)
# Convert to tensor without resizing or additional transformations
inputs = self.vision_encoder.image_processor(images=[image], return_tensors="pt")
pixel_values = inputs.pixel_values.to(self.param_dtype).to(self.device)
return pixel_values
def extract_embedding(self, image_tensor):
"""Extract raw vision embedding."""
with torch.no_grad():
vision_output = self.vision_encoder(image_tensor)
vision_output = vision_output.mean(axis=0)
return vision_output
def save_embedding(self, vision_output, file_path):
"""Save the vision output to a numpy file."""
np.save(file_path, vision_output.cpu().numpy())
def process_directory(self, image_paths, output_dir):
"""Process all images in a directory with a progress bar and save the embeddings."""
if os.path.exists(output_dir):
shutil.rmtree(output_dir)
print(f"Existing directory cleared: {output_dir}")
os.makedirs(output_dir, exist_ok=True)
# Adding tqdm for progress bar
for image_path in tqdm(image_paths, desc="Processing Images", unit="image"):
# Load and extract features without preprocessing
image_tensor = self.load_image(image_path)
vision_output = self.extract_embedding(image_tensor)
# Save the output with the same filename but as a .npy
output_file_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(image_path))[0]}.npy")
self.save_embedding(vision_output, output_file_path)
if __name__ == "__main__":
torch.manual_seed(42)
device = "cuda" if torch.cuda.is_available() else "cpu"
param_dtype = torch.float32
# Instantiate the pipeline
pipeline = VisionPreprocessor(device, param_dtype)
# Specify input and output directories
input_directory = "/mnt/nvme/shared_A/datasets/flickr30k/data/flickr30k-images"
output_directory = "/mnt/nvme/shared_A/datasets/flickr30k/data/vision_embeddings_reduced2"
image_paths = [os.path.join(input_directory, f) for f in os.listdir(input_directory) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
# Process all images in the input directory
pipeline.process_directory(image_paths, output_directory)
print("Processing complete!")
|