Fix app file

app.py

@@ -1,204 +1,203 @@
-import spaces
-import torch
-import torch.nn as nn
-from huggingface_hub import hf_hub_download
-import librosa
-import numpy as np
+import cv2
 import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+import torch
+import open_clip
+import chromadb
+
+# 1. Device and Model Initialization
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, _, preprocess = open_clip.create_model_and_transforms(
+    "ViT-B-32",
+    pretrained="laion2b_s34b_b79k",
+    device=device
+)
+tokenizer = open_clip.get_tokenizer("ViT-B-32")
 
+# 2. Visual Attribution Heatmap Generation Function
+def generate_visual_attribution(model, preprocess, rgb_image, query_string, grid_size=4):
+    """Generates a side-by-side visualization array with a semantic heatmap overlay."""
+    h, w, _ = rgb_image.shape
+    patch_h, patch_w = h // grid_size, w // grid_size
+    heatmap = np.zeros((grid_size, grid_size))
 
-# Vars
-repo_name = "ipetrousov/ser_supcon"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-weights_dict = "supcon_stage_ft_acc60.pt"
-
-# Classes
-class EmotionFeatureExtractorBase(nn.Module):
-    """
-    Core CNN to extract spatial features from Mel-Spectrograms.
-    Outputs a flat 256-dimensional feature vector.
-    """
-
-    def __init__(self):
-        super(EmotionFeatureExtractorBase, self).__init__()
-
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(32),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2)
-        )
-
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(64),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2)
-        )
-
-        self.conv3 = nn.Sequential(
-            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(128),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2)
-        )
-
-        self.conv4 = nn.Sequential(
-            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
-            nn.BatchNorm2d(256),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2)
-        )
-
-        self.global_pool = nn.AdaptiveAvgPool2d((1, 1))
-        self.flatten = nn.Flatten()
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = self.conv3(x)
-        x = self.conv4(x)
-        x = self.global_pool(x)
-        x = self.flatten(x)
-        return x # Output shape: (batch_size, 256)
-
-
-class SupConStage2(nn.Module):
-    """ Model B (Stage 2): Frozen Encoder + Linear Classifier. """
-
-    def __init__(self, trained_encoder, num_classes=8):
-        super(SupConStage2, self).__init__()
-
-        # Pass trained encoder from Stage 1
-        self.encoder = trained_encoder
-
-        # Freeze encoder weights
-        for param in self.encoder.parameters():
-            param.requires_grad = False
-
-        # Linear classifier
-        self.classifier = nn.Linear(256, num_classes)
-
-    def forward(self, x):
-        # Disable gradient tracking for the encoder to save memory/compute
-        with torch.no_grad():
-            features = self.encoder(x)
-
-        return self.classifier(features)
-
-
-# Function
-def preprocess_drive_audio(file_path, target_sample_rate=22050, n_mels=128, noise_std=0.0):
-    """
-    Loads audio and generates a normalized Mel-spectrogram using librosa.
-    Returns a PyTorch tensor shaped (1, 1, Mels, Time).
-    """
-
-    y, sr = librosa.load(file_path, sr=target_sample_rate, mono=True)
-
-    mel_spectrogram = librosa.feature.melspectrogram(
-        y=y,
-        sr=target_sample_rate,
-        n_fft=512,
-        hop_length=256,
-        n_mels=n_mels,
-        fmax=8000
+    # Tokenize and vectorize text query
+    text_token = tokenizer([query_string]).to(device)
+    with torch.no_grad():
+        text_features = model.encode_text(text_token)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+
+    # Calculate patch-level cosine similarities
+    for i in range(grid_size):
+        for j in range(grid_size):
+            ymin, ymax = i * patch_h, (i + 1) * patch_h
+            xmin, xmax = j * patch_w, (j + 1) * patch_w
+
+            patch = rgb_image[ymin:ymax, xmin:xmax]
+            pil_patch = Image.fromarray(patch)
+            patch_tensor = preprocess(pil_patch).unsqueeze(0).to(device)
+
+            with torch.no_grad():
+                patch_features = model.encode_image(patch_tensor)
+                patch_features /= patch_features.norm(dim=-1, keepdim=True)
+                similarity = (patch_features @ text_features.T).item()
+                heatmap[i, j] = similarity
+
+    heatmap_resized = cv2.resize(heatmap, (w, h), interpolation=cv2.INTER_CUBIC)
+
+    # Render side-by-side figure
+    fig, axes = plt.subplots(1, 2, figsize=(12, 6))
+    axes[0].imshow(rgb_image)
+    axes[0].set_title("Original Matched Frame")
+    axes[0].axis("off")
+
+    axes[1].imshow(rgb_image)
+    axes[1].imshow(heatmap_resized, cmap="jet", alpha=0.5)
+    axes[1].set_title(f"Visual Attribution Heatmap for:\n\"{query_string}\"")
+    axes[1].axis("off")
+
+    plt.tight_layout()
+
+    # Convert Matplotlib figure directly to an RGB array for Gradio
+    fig.canvas.draw()
+    output_image = np.asarray(fig.canvas.buffer_rgba())[:, :, :3]
+    plt.close(fig)
+
+    return output_image
+
+# 3. Main Pipeline Execution
+def process_video_and_query(video_path, query_string, sample_rate_sec=1):
+    if not video_path:
+        return None, "Please upload a video file."
+    if not query_string:
+        return None, "Please enter a search query."
+
+    # Initialize fresh in-memory ChromaDB collection per query
+    chroma_client = chromadb.Client()
+    collection = chroma_client.create_collection(name="temp_video_analysis")
+
+    # Ingest and sample video frames
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    if fps == 0 or np.isnan(fps):
+        fps = 30.0
+    frame_interval = int(fps * sample_rate_sec)
+
+    frame_count = 0
+    saved_frames = {}
+    embeddings = []
+    ids = []
+    metadatas = []
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        if frame_count % frame_interval == 0:
+            timestamp = round(frame_count / fps, 2)
+            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            saved_frames[str(timestamp)] = rgb_frame
+
+            pil_img = Image.fromarray(rgb_frame)
+            img_tensor = preprocess(pil_img).unsqueeze(0).to(device)
+
+            with torch.no_grad():
+                img_embedding = model.encode_image(img_tensor).flatten().tolist()
+
+            embeddings.append(img_embedding)
+            ids.append(f"frame_{timestamp}")
+            metadatas.append({"timestamp": str(timestamp)})
+
+        frame_count += 1
+    cap.release()
+
+    if not embeddings:
+        return None, "Error: Could not extract frames from the video."
+
+    # Batch add to ChromaDB
+    collection.add(
+        embeddings=embeddings,
+        ids=ids,
+        metadatas=metadatas
     )
 
-    mel_spectrogram_db = librosa.power_to_db(mel_spectrogram, ref=1.0)
-
-    mean = np.mean(mel_spectrogram_db)
-    std = np.std(mel_spectrogram_db)
-    mel_normalized = (mel_spectrogram_db - mean) / (std + 1e-6)
-
-    final_tensor = torch.tensor(mel_normalized, dtype=torch.float32)
-    final_tensor = final_tensor.unsqueeze(0).unsqueeze(0)
-
-    # Noise Injection Pipeline
-    if noise_std > 0.0:
-        noise_matrix = torch.randn_like(final_tensor)
-        final_tensor = final_tensor + (noise_matrix * noise_std)
-
-    return final_tensor
+    # Vectorize text query
+    text_token = tokenizer([query_string]).to(device)
+    with torch.no_grad():
+        query_embedding = model.encode_text(text_token).flatten().tolist()
 
+    # Retrieve nearest neighbor
+    results = collection.query(
+        query_embeddings=[query_embedding],
+        n_results=1
+    )
 
-def run_local_inference(model, file_path, device, noise_std=0.0):
-    """
-    Runs a forward pass
-    """
+    best_timestamp = results["metadatas"][0][0]["timestamp"]
+    distance = results["distances"][0][0]
+    matched_frame = saved_frames[best_timestamp]
 
-    emotion_classes = ["Neutral", "Calm", "Happy", "Sad", "Angry", "Fearful", "Disgust", "Surprised"]
+    # Generate side-by-side attribution image
+    side_by_side_output = generate_visual_attribution(
+        model=model,
+        preprocess=preprocess,
+        rgb_image=matched_frame,
+        query_string=query_string
+    )
 
-    input_tensor = preprocess_drive_audio(file_path, noise_std=noise_std)
-    input_tensor = input_tensor.to(device)
+    status_message = (
+        f"Matched Timestamp: {best_timestamp}s\n"
+        f"ChromaDB Squared Euclidean Distance: {distance:.4f}"
+    )
 
-    model.eval()
-    with torch.no_grad():
-        logits = model(input_tensor)
-        
-        # Get the top prediction index
-        predicted_index = torch.argmax(logits, dim=1).item()
-        top_class = emotion_classes[predicted_index]
-        
-        # Apply softmax to convert logits to probabilities between 0 and 1
-        probabilities = torch.softmax(logits, dim=1).squeeze().tolist()
-
-    # Map the probabilities to their respective class names for Gradio
-    prob_dict = {emotion_classes[i]: probabilities[i] for i in range(len(emotion_classes))}
-    
-    return top_class, prob_dict
-
-
-# Load SOTA SupConFN
-blank_encoder = EmotionFeatureExtractorBase()
-model = SupConStage2(trained_encoder=blank_encoder, num_classes=8)
-
-# Download & Load weights
-model_path = hf_hub_download(repo_id=repo_name, filename=weights_dict)
-model.load_state_dict(torch.load(model_path, map_location=device))
-model.eval()
-model = model.to(device)
-
-
-# Gradio Interface Setup
-@spaces.GPU
-def predict_emotion(audio_path, noise_std):
-    if audio_path is None:
-        return "Upload an audio file."
-    
-    try:
-        pred = run_local_inference(model, audio_path, device, noise_std=noise_std)
-        return pred
-    except Exception as e:
-        return f"Error processing audio: {str(e)}"
+    return side_by_side_output, status_message
 
 
 description_html = """
-Upload a short audio clip and test the architecture's resilience. The slider injects mathematically controlled Gaussian noise into the normalized Mel-spectrogram before classification.
+Upload surveillance footage and enter a natural language safety policy to retrieve candidate violation frames via Open-CLIP embeddings and ChromaDB nearest-neighbor search, complete with patch-level visual attribution heatmaps.
 
-<div style="display: flex; gap: 10px; margin-top: 15px;">
-    <a href="https://github.com/gpetrousov/dl_assignment_demokritos" target="_blank">
+<div style="display: flex; gap: 10px; margin-top: 15px; flex-wrap: wrap;">
+    <a href="https://github.com/gpetrousov/multimodal_ai_assignment_demokritos" target="_blank">
         <img src="https://img.shields.io/badge/GitHub-View_Repository-181717?style=for-the-badge&logo=github" alt="GitHub Repository" />
     </a>
-    <a href="https://docs.google.com/presentation/d/1ShSmXY2OqCEbQQjYiLYCFc8Psu2bioECoFHhliiDEiQ/edit?usp=drive_link" target="_blank">
+    <a href="https://github.com/gpetrousov/multimodal_ai_assignment_demokritos/blob/master/report_src/report.pdf" target="_blank">
+        <img src="https://img.shields.io/badge/Report-View_PDF-E50914?style=for-the-badge&logo=adobeacrobatreader&logoColor=white" alt="View PDF Report" />
+    </a>
+    <a href="https://docs.google.com/presentation/d/1DW28UWfK28nR-G_rB0mqrZmNZY3uNZMmDJpCxxdA-fk/edit?usp=sharing" target="_blank">
         <img src="https://img.shields.io/badge/Presentation-View_Slides-F4B400?style=for-the-badge&logo=googleslides&logoColor=white" alt="Presentation Slides" />
     </a>
 </div>
 """
 
-demo = gr.Interface(
-    fn=predict_emotion,
-    inputs=[
-        gr.Audio(type="filepath", label="Upload Audio (.mp3 or .wav)"),
-        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.0, label="Gaussian Noise Injection (STD)")
-    ],
-    outputs=[
-        gr.Textbox(label="Final Prediction"),
-        gr.Label(num_top_classes=8, label="Probability Distribution")
-    ],
-    title="Robust Speech Emotion Recognition (SupCon)",
-    description=description_html,
-)
+# 4. Gradio Interface Layout
+with gr.Blocks(title="Zero-Shot Video Safety Auditor") as demo:
+    gr.Markdown("# Zero-Shot Video Safety Policy Auditor")
+    gr.Markdown(
+        "Upload surveillance footage and enter a natural language safety policy "
+        "to retrieve candidate violation frames with visual attribution heatmaps."
+    )
+    gr.HTML(description_html)
+
+    with gr.Row():
+        with gr.Column():
+            video_input = gr.Video(label="Input Surveillance Video")
+            query_input = gr.Textbox(
+                label="Safety Policy Query",
+                placeholder="e.g., person operating a yellow construction machine"
+            )
+            submit_btn = gr.Button("Analyze & Retrieve", variant="primary")
+
+        with gr.Column():
+            image_output = gr.Image(label="Side-by-Side Retrieval & Attribution Heatmap")
+            details_output = gr.Textbox(label="Retrieval Metadata", interactive=False)
+
+    submit_btn.click(
+        fn=process_video_and_query,
+        inputs=[video_input, query_input],
+        outputs=[image_output, details_output]
+    )
 
 if __name__ == "__main__":
     demo.launch()