CLIP Model

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
    

def create_image_encoder(input_shape = (224, 224, 3)):
    model = tf.keras.Sequential([
        tf.keras.layers.InputLayer(input_shape = input_shape),
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same'),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(512, activation='relu'),
        tf.keras.layers.Dense(256)  # Final feature size for image
    ])
    return model

def create_text_encoder(vocab_size=10000, embedding_dim=256, seq_length=32):
    inputs = tf.keras.Input(shape=(seq_length,), dtype=tf.int32)
    embedding_layer = tf.keras.layers.Embedding(vocab_size, embedding_dim)(inputs)
    transformer_layer = tf.keras.layers.MultiHeadAttention(num_heads=4, key_dim=embedding_dim)(embedding_layer, embedding_layer)
    transformer_layer = tf.keras.layers.GlobalAveragePooling1D()(transformer_layer) 
    dense_layer = tf.keras.layers.Dense(256, activation='relu')(transformer_layer)
    output = tf.keras.layers.Dense(256)(dense_layer)  # Final feature size for text
    model = tf.keras.Model(inputs=inputs, outputs=output)
    return model

def contrastive_loss(image_embeds, text_embeds, temperature=0.07):
    # Normalize embeddings to unit vectors
    image_embeds = tf.nn.l2_normalize(image_embeds, axis=-1)
    text_embeds = tf.nn.l2_normalize(text_embeds, axis=-1)    
    # Compute cosine similarity
    similarity = tf.matmul(image_embeds, text_embeds, transpose_b=True)
    # Scale similarity by temperature
    similarity /= temperature
    # Create labels (same image-text pair should have label 1)
    labels = tf.eye(image_embeds.shape[0])
    # Compute the loss using softmax
    loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=similarity)
    return tf.reduce_mean(loss)

@tf.function
def train_step(image_batch, text_batch, image_encoder, text_encoder, optimizer):
    with tf.GradientTape() as tape:
        # Extract embeddings for both image and text
        image_embeds = image_encoder(image_batch)  # Shape: (batch_size, embed_dim)
        text_embeds = text_encoder(text_batch)    # Shape: (batch_size, embed_dim)
        # Calculate contrastive loss
        loss = contrastive_loss(image_embeds, text_embeds)
    # Compute gradients and update weights
    gradients = tape.gradient(loss, image_encoder.trainable_variables + text_encoder.trainable_variables)
    optimizer.apply_gradients(zip(gradients, image_encoder.trainable_variables + text_encoder.trainable_variables))
    return loss

def preprocess_image(image_path):
    image = image = tf.io.read_file(join('/kaggle/input/clipimages/images/images', image_path))
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, (224, 224))  # Resize to 224x224
    image = image / 255.0  # Normalize to [0, 1]
    return image
def preprocess_text(text, tokenizer):
    return tokenizer(text)

# Example image paths and text descriptions
image_paths = df['image'].values[:1000]
texts = df['caption'].values[:1000]
# Define a simple tokenizer for text processing (you can replace this with a more advanced one)
max_len = 32
tokenizer = tf.keras.layers.TextVectorization(max_tokens=10000, output_sequence_length=max_len)
tokenizer.adapt(texts)

# Instantiate image and text encoders
image_encoder = create_image_encoder()
text_encoder = create_text_encoder()
# Optimizer
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
# Training loop
batch_size = 200  # Adjust based on your dataset size
num_epochs = 100
for epoch in range(num_epochs):
    for i in range(0, len(image_paths), batch_size):
        # Load and preprocess images and texts
        batch_images = [preprocess_image(p) for p in image_paths[i:i+batch_size]]
        batch_texts = tokenizer(texts[i:i+batch_size])
        batch_images = tf.stack(batch_images)
        batch_texts = tf.stack(batch_texts)
        # Train the model
        loss = train_step(batch_images, batch_texts, image_encoder, text_encoder, optimizer)
        if i % 100 == 0:  # Print loss every 100 steps
            print(f"Epoch {epoch}, Step {i}, Loss: {loss.numpy()}")