[NEW] Sound modality input. Yet to put in proper examples

app.py

@@ -10,10 +10,13 @@ import numpy as np
 from PIL import Image
 import matplotlib.pyplot as plt
 import io
+import torchaudio
 
 from torchvision import transforms
 import open_clip
 from clip_vision_per_patch_model import CLIPVisionPerPatchModel
+from transformers import ClapAudioModelWithProjection
+from transformers import ClapProcessor
 
 # ────────────────────────── global config & models ────────────────────
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -25,11 +28,20 @@ bio_tokenizer = open_clip.get_tokenizer("hf-hub:imageomics/bioclip")
 
 # Satellite patch encoder CLIP-L-336 per-patch)
 sat_model: CLIPVisionPerPatchModel = (
-    CLIPVisionPerPatchModel.from_pretrained("derektan95/search-tta")
+    CLIPVisionPerPatchModel.from_pretrained("derektan95/search-tta-sat")
     .to(device)
     .eval()
 )
 
+# Sound CLAP model
+sound_model: ClapAudioModelWithProjection = (
+    ClapAudioModelWithProjection.from_pretrained("derektan95/search-tta-sound")
+    .to(device)
+    .eval()
+)
+sound_processor: ClapProcessor = ClapProcessor.from_pretrained("derektan95/search-tta-sound")
+SAMPLE_RATE = 48000
+
 logit_scale = torch.nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
 logit_scale = logit_scale.exp()
 blur_kernel = (5,5)
@@ -58,6 +70,13 @@ imo_transform = transforms.Compose(
     ]
 )
 
+def get_audio_clap(path_to_audio,format="mp3",padding="repeatpad",truncation="fusion"):
+    track, sr = torchaudio.load(path_to_audio, format=format)  # torchaudio.load(path_to_audio)
+    track = track.mean(axis=0)
+    track = torchaudio.functional.resample(track, orig_freq=sr, new_freq=SAMPLE_RATE)
+    output = sound_processor(audios=track, sampling_rate=SAMPLE_RATE, max_length_s=10, return_tensors="pt",padding=padding,truncation=truncation)
+    return output
+
 # ────────────────────────── helpers ───────────────────────────────────
 
 @torch.no_grad()
@@ -81,6 +100,16 @@ def _encode_sat(img_pil: Image.Image) -> torch.Tensor:
     return imo_embeds
 
 
+@torch.no_grad()
+def _encode_sound(sound) -> torch.Tensor:
+    processed_sound = get_audio_clap(sound)
+    for k in processed_sound.keys():
+        processed_sound[k] = processed_sound[k].to(device)
+    unnormalized_audio_embeds = sound_model(**processed_sound).audio_embeds
+    sound_embeds = torch.nn.functional.normalize(unnormalized_audio_embeds, dim=-1)
+    return sound_embeds
+
+
 def _similarity_heatmap(query: torch.Tensor, patches: torch.Tensor) -> np.ndarray:
     sims = torch.matmul(query, patches.t()) * logit_scale
     sims = sims.t().sigmoid()
@@ -122,13 +151,14 @@ def process(
     sat_img: Image.Image,
     taxonomy: str,
     ground_img: Image.Image | None,
+    sound: torch.Tensor | None,
 ):
     if sat_img is None:
         return None, None
 
     patches = _encode_sat(sat_img)
 
-    heat_ground, heat_text = None, None
+    heat_ground, heat_text, heat_sound = None, None, None
 
     if ground_img is not None:
         q_img = _encode_ground(ground_img)
@@ -138,7 +168,11 @@ def process(
         q_txt = _encode_text(taxonomy.strip())
         heat_text = _array_to_pil(_similarity_heatmap(q_txt, patches))
 
-    return heat_ground, heat_text
+    if sound is not None:
+        q_sound = _encode_sound(sound)
+        heat_sound = _array_to_pil(_similarity_heatmap(q_sound, patches))
+
+    return heat_ground, heat_text, heat_sound
 
 
 # ────────────────────────── Gradio UI ─────────────────────────────────
@@ -191,6 +225,13 @@ with gr.Blocks(title="Search-TTA", theme=gr.themes.Base()) as demo:
                 label="Full Taxonomy Name (optional)",
                 placeholder="e.g. Animalia Chordata Mammalia Rodentia Sciuridae Marmota marmota",
             )
+
+            # ─── NEW: sound input ───────────────────────────
+            sound_input = gr.Audio(
+                label="Sound Input",
+                source="upload",     # or "microphone" / "url" as you prefer
+                type="filepath",     # or "numpy" if you want raw arrays
+            )
             run_btn = gr.Button("Run", variant="primary")
 
         # RIGHT COLUMN  (ground image + two heat-maps)
@@ -209,6 +250,15 @@ with gr.Blocks(title="Search-TTA", theme=gr.themes.Base()) as demo:
                 label="Heatmap (Text query)",
                 height=160,
             )
+            heat_sound_out = gr.Image(
+                label="Heatmap (Sound query)",
+                height=160,
+            )
+            # ─── NEW: sound output ─────────────────────────
+            # sound_output = gr.Audio(
+            #     label="Playback",
+            # )
+            
 
     # EXAMPLES
     with gr.Row():
@@ -218,25 +268,29 @@ with gr.Blocks(title="Search-TTA", theme=gr.themes.Base()) as demo:
                     "examples/Animalia_Chordata_Mammalia_Carnivora_Ursidae_Ursus_americanus_NAIP_yosemite_v3_resized.png",
                     "Animalia Chordata Mammalia Carnivora Ursidae Ursus americanus",
                     "examples/Animalia_Chordata_Mammalia_Carnivora_Ursidae_Ursus_americanus_inat_248820933.jpeg",
+                    None
                 ],
                 [
                     "examples/Animalia_Chordata_Reptilia_Squamata_Varanidae_Varanus_salvator_sentinel2_410613_5.35573_100.28948.jpg",
                     "Animalia Chordata Reptilia Squamata Varanidae Varanus salvator",
                     "examples/Animalia_Chordata_Reptilia_Squamata_Varanidae_Varanus_salvator_inat_461d8e6c-0e66-4acc-8ecd-bfd9c218bc14.jpg",
+                    None
+                ],
+                [
+                    "examples/Animalia_Chordata_Mammalia_Rodentia_Sciuridae_Marmota_marmota_sentinel2_388246_45.49036_7.14796.jpg",
+                    "Animalia Chordata Mammalia Rodentia Sciuridae Marmota marmota",
+                    "examples/Animalia_Chordata_Mammalia_Rodentia_Sciuridae_Marmota_marmota_inat_327e1f07-692b-4140-8a3e-bd098bc064ff.jpg",
+                    None
                 ],
                 [
                     "examples/Animalia_Chordata_Mammalia_Rodentia_Sciuridae_Marmota_marmota_sentinel2_388246_45.49036_7.14796.jpg",
                     "Animalia Chordata Mammalia Rodentia Sciuridae Marmota marmota",
                     "examples/Animalia_Chordata_Mammalia_Rodentia_Sciuridae_Marmota_marmota_inat_327e1f07-692b-4140-8a3e-bd098bc064ff.jpg",
+                    "/mnt/hdd/inat2021_ds/2_OTHERS/sound_test/sounds_mp3/386157.mp3"
                 ],
-                # [
-                #     "examples/satellite_coast.png",
-                #     "Animalia Chordata Aves Charadriiformes Laridae Larus argentatus",
-                #     None,
-                # ],
             ],
-            inputs=[sat_input, taxonomy_input, ground_input],
-            outputs=[heat_ground_out, heat_text_out],
+            inputs=[sat_input, taxonomy_input, ground_input, sound_input],
+            outputs=[heat_ground_out, heat_text_out, heat_sound_out],
             fn=process,
             cache_examples=False,
         )
@@ -244,8 +298,8 @@ with gr.Blocks(title="Search-TTA", theme=gr.themes.Base()) as demo:
     # CALLBACK
     run_btn.click(
         fn=process,
-        inputs=[sat_input, taxonomy_input, ground_input],
-        outputs=[heat_ground_out, heat_text_out],
+        inputs=[sat_input, taxonomy_input, ground_input, sound_input],
+        outputs=[heat_ground_out, heat_text_out, heat_sound_out],
     )
 
     # Footer to point out to model and data from app page.