Add annotations

app.py

@@ -0,0 +1,98 @@
+# Construct pairs of text and image
+from configs import CFG
+from costum_datasets import make_pairs
+
+
+from text_image_audio import OneEncoder
+import torch
+
+import gradio as gr
+
+import torchaudio
+
+# Construct pairs of text and image
+training_pairs = make_pairs(CFG.train_annotation_file, CFG.image_dir, 5) # 413.915 -> 82.783 images
+
+# Sorted according images
+training_pairs = sorted(training_pairs, key=lambda x: x[0])
+
+coco_images, coco_captions = zip(*training_pairs)
+
+# Take unique images
+unique_images = set()
+unique_pairs = [(item[0], item[1]) for item in training_pairs if item[0] not in unique_images
+                and not unique_images.add(item[0])]
+coco_images, _ = zip(*unique_pairs)
+
+# Load model
+model = OneEncoder.from_pretrained("bilalfaye/OneEncoder-text-image-audio")
+
+# Load coco image features
+coco_image_features = torch.load("image_embeddings_best.pt", map_location=CFG.device)
+coco_image_features = coco_image_features[:3000]
+
+def text_image(query):
+    model.text_image_encoder.image_retrieval(query,
+                          image_paths=coco_images,
+                          image_embeddings=coco_image_features,
+                          n=9,
+                          plot=True,
+                          temperature=0.0
+                          )
+    return "img.png"
+
+def audio_image(query):
+    # Load the audio with torchaudio (returns tensor and sample rate)
+    waveform, sample_rate = torchaudio.load(query)
+
+    # Check if audio is stereo
+    if waveform.shape[0] > 1:  # Stereo (2 channels)
+        # Convert stereo to mono: sum the left and right channels and divide by 2
+        mono_audio = waveform.mean(dim=0, keepdim=True)
+    else:
+        # Audio is already mono
+        mono_audio = waveform
+
+    # Resample to 16000 Hz if not already
+    if sample_rate != 16000:
+        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
+        mono_audio = resampler(mono_audio)
+
+    # Convert to numpy array for pipeline processing (if required)
+    mono_audio = mono_audio.squeeze(0).numpy()
+
+    audio_encoding = model.process_audio([mono_audio])
+
+    model.image_retrieval(audio_encoding, 
+                          image_paths=coco_images,
+                          image_embeddings=coco_image_features, 
+                          n=9, 
+                          plot=True, 
+                          temperature=0.0,
+                          display_audio=False)
+
+    return "img.png"
+
+
+# Updated Gradio Interface
+iface = gr.TabbedInterface(
+    [
+        gr.Interface(
+            fn=text_image,
+            inputs=gr.Textbox(label="Text Query"),
+            outputs="image",
+            title="Retrieve images using text as query",
+            description="Implementation of OneEncoder using one layer on UP for light demo, Only coco train dataset is used in this example (3000 images)."
+        ),
+        gr.Interface(
+            fn=audio_image,
+            inputs=gr.Audio(sources=["upload", "microphone"], type="filepath", label="Provide Audio Query"),
+            outputs="image",
+            title="Retrieve images using audio as query",
+            description="Implementation of OneEncoder using one layer on UP for light demo, Only coco train dataset is used in this example (3000 images)."
+        )
+    ],
+    tab_names=["Text - Image", "Audio - Image"]
+)
+
+iface.launch(debug=True, share=True)

configs.py

@@ -0,0 +1,40 @@
+import torch
+import os
+
+################################################### PARMETERS ##########################################################
+################################################# PARAMETERS ###########################################################
+
+class CFG:
+    max_length = 128
+    batch_size = 32
+    num_workers = 4
+    projection_dim = 768
+    dropout_rate = 0.1
+    num_head = 4
+    num_layers = 1
+    image_encoder_lr = 1e-4
+    radio_encoder_lr = 1e-5
+    video_encoder_lr = 1e-4
+    text_encoder_lr = 1e-5
+    audio_encoder_lr = 1e-5
+    modality_token_encoder_lr = 1e-3
+    universal_projection_lr = 1e-3
+    lr = 1e-3
+    weight_decay = 1e-3
+    patience = 10
+    factor = 0.8
+    token_size = 1
+    epochs = 100
+    image_size = 224
+    device = "cpu"
+    data_directory = "datasets"
+    train_annotation_file = os.path.join(data_directory, "annotations", "captions_train2014.json")
+    val_annotation_file = os.path.join(data_directory, "annotations", "captions_val2014.json")
+    image_dir = os.path.join(data_directory, "train2014")
+    image_dir_val = os.path.join(data_directory, "val2014")
+    bert_name = "bert-base-uncased"
+    vit_name = "vit_base_patch16_224"
+    audio_name = "facebook/wav2vec2-base-960h"
+    radio_name = "microsoft/rad-dino"
+    video_name = "MCG-NJU/videomae-base"
+    sample_rate = 16000

costum_datasets.py

@@ -0,0 +1,67 @@
+from configs import CFG
+import os
+import requests
+import zipfile
+from pycocotools.coco import COCO
+import torch
+import cv2
+import albumentations as A
+import soundfile as sf
+
+
+
+# Load Coco dataset
+def download_dataset(data_dir="../datasets"):
+    # Create caption and image directories
+    annotations_dir = os.path.join(data_dir, "annotations")
+    images_dir = os.path.join(data_dir, "train2014")
+
+    # Download annotations (captions)
+    zip_file = os.path.join(annotations_dir, "annotations.zip")
+    url = "http://images.cocodataset.org/annotations/annotations_trainval2014.zip"
+    response = requests.get(url, stream=True)
+    # write chunk in zip file
+    with open(zip_file, "wb") as f:
+        # 8192 = 8KB chunks (block or piece of data)
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+    # unzip file
+    with zipfile.ZipFile(zip_file, "r") as zip_ref:
+        zip_ref.extractall(data_dir)  # Extract all contents to the specified directory
+    os.remove(zip_file)
+
+    # Download train images
+    zip_file = os.path.join(images_dir, "train2014.zip")
+    url = "http://images.cocodataset.org/zips/train2014.zip"
+    response = requests.get(url, stream=True)
+    # write chunk in zip file
+    with open(zip_file, "wb") as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+    # unzip file
+    with zipfile.ZipFile(zip_file, "r") as zip_ref:
+        zip_ref.extractall(data_dir)  # Extract all contents to the specified directory
+    os.remove(zip_file)
+
+    # Download val images
+    images_dir = os.path.join(data_dir, "val2014")
+    zip_file = os.path.join(images_dir, "val2014.zip")
+    url = "http://images.cocodataset.org/zips/val2014.zip"
+    response = requests.get(url, stream=True)
+    # write chunk in zip file
+    with open(zip_file, "wb") as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+    # unzip file
+    with zipfile.ZipFile(zip_file, "r") as zip_ref:
+        zip_ref.extractall(data_dir)  # Extract all contents to the specified directory
+    os.remove(zip_file)
+
+
+
+def make_pairs(annotation_json_files, image_dir, max_captions=3):
+    
+    images = os.listdir(annotation_json_files)
+    image_caption = [(os.path.join(annotation_json_files, image), "an image") for image in images]
+    return image_caption
+