lambda_function.py

import os
import json
import time
import uuid
import boto3
from PIL import Image
import numpy as np
from io import BytesIO
from concurrent.futures import ThreadPoolExecutor, as_completed
from supabase import create_client, Client
from dotenv import load_dotenv
import torch
import clip

load_dotenv()

# Initialize S3 client and bucket name from environment variable
s3_client = boto3.client('s3')
BUCKET_NAME = os.environ.get("BUCKET_NAME", "oriane-contents")

# Initialize Supabase client using environment variables
SUPABASE_URL = os.environ.get("SUPABASE_URL")
SUPABASE_KEY = os.environ.get("SUPABASE_KEY")
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

# Load CLIP model and preprocessing function
device = "cuda" if torch.cuda.is_available() else "cpu"
# Use the environment variable if available, else fall back to /tmp/clip
cache_dir = os.environ.get("CLIP_DOWNLOAD_ROOT", "/tmp/clip")
os.makedirs(cache_dir, exist_ok=True)
model, preprocess = clip.load("ViT-B/32", device=device, download_root=cache_dir)

def download_frame(shortcode, frame_number, platform, extension):
    """Download a single frame from S3."""
    key = f"{platform}/{shortcode}/frames/{frame_number}.{extension}"
    try:
        response = s3_client.get_object(Bucket=BUCKET_NAME, Key=key)
        return Image.open(BytesIO(response['Body'].read()))
    except s3_client.exceptions.NoSuchKey:
        return None

def get_all_frames(shortcode, platform, extension):
    """
    Get all available frames for a given shortcode.
    Uses a ThreadPoolExecutor to download frames concurrently.
    Assumes a maximum of 100 frames; adjust as needed.
    """
    frames = []
    max_frames = 100  # adjust if necessary
    with ThreadPoolExecutor() as executor:
        future_to_frame = {
            executor.submit(download_frame, shortcode, i, platform, extension): i 
            for i in range(max_frames)
        }
        for future in as_completed(future_to_frame):
            frame_number = future_to_frame[future]
            frame = future.result()
            if frame:
                frames.append((frame_number, frame))
        # Sort frames by frame number to maintain order
        frames.sort(key=lambda x: x[0])
    return [frame for _, frame in frames]

def extract_features(image: Image.Image):
    """
    Convert a PIL image to a feature vector using CLIP.
    Preprocess the image, encode it with the CLIP model,
    and return a feature vector.
    """
    image_input = preprocess(image).unsqueeze(0).to(device)
    with torch.no_grad():
        features = model.encode_image(image_input)
    return features.cpu().numpy().squeeze()

def cosine_similarity(a, b):
    """Compute the cosine similarity between two vectors."""
    a = np.array(a)
    b = np.array(b)
    dot_product = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot_product / (norm_a * norm_b)

def compare_frames(frame1, frame2):
    """
    Compare two frames using deep features from CLIP.
    Extract features from both frames and compute the cosine similarity.
    """
    feat1 = extract_features(frame1)
    feat2 = extract_features(frame2)
    similarity = cosine_similarity(feat1, feat2)
    return similarity

def lambda_handler(event, context):
    try:
        # Check if a job_id is provided; if not, generate one and insert into ai_jobs.
        job_id = event.get('job_id')
        if not job_id:
            job_id = str(uuid.uuid4())
            job_insert_response = supabase.table("ai_jobs").insert({"job_id": job_id}).execute()
            if job_insert_response.error:
                raise Exception(f"Error inserting job: {job_insert_response.error}")

        # Extract parameters from the event
        monitored_shortcode = event.get('monitored_shortcode')
        watched_shortcodes = event.get('watched_shortcodes', [])
        platform = event.get('platform', 'instagram')
        extension = event.get('extension', 'jpg')
        
        if not monitored_shortcode or not watched_shortcodes:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Missing required parameters: monitored_shortcode and watched_shortcodes'})
            }
        
        # Download frames for the monitored video
        monitored_frames = get_all_frames(monitored_shortcode, platform, extension)
        if not monitored_frames:
            return {
                'statusCode': 404,
                'body': json.dumps({'error': f'No frames found for monitored shortcode: {monitored_shortcode}'})
            }
        
        records_to_insert = []
        
        # Process each watched video
        for watched_shortcode in watched_shortcodes:
            start_time_video = time.time()
            watched_frames = get_all_frames(watched_shortcode, platform, extension)
            
            if not watched_frames:
                record = {
                    "job_id": job_id,
                    "monitored_video": monitored_shortcode,
                    "watched_video": watched_shortcode,
                    "avg_similarity": None,
                    "processed_in_secs": time.time() - start_time_video,
                    "frame_results": [],
                    "max_similarity": None
                }
                records_to_insert.append(record)
                continue
            
            frame_comparisons = []
            # Compare frames one-by-one
            for i, (monitored_frame, watched_frame) in enumerate(zip(monitored_frames, watched_frames)):
                similarity = compare_frames(monitored_frame, watched_frame)
                frame_comparisons.append({'frame_number': i, 'similarity': float(similarity)})
            
            similarities = [comp['similarity'] for comp in frame_comparisons]
            avg_similarity = float(np.mean(similarities)) if similarities else None
            max_similarity = float(max(similarities)) if similarities else None
            processed_time = time.time() - start_time_video
            
            record = {
                "job_id": job_id,
                "monitored_video": monitored_shortcode,
                "watched_video": watched_shortcode,
                "avg_similarity": avg_similarity,
                "processed_in_secs": processed_time,
                "frame_results": frame_comparisons,
                "max_similarity": max_similarity
            }
            records_to_insert.append(record)
        
        # Bulk insert the results into the ai_results table in Supabase.
        supabase_response = supabase.table("ai_results").insert(records_to_insert).execute()
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Analysis complete and stored in Supabase',
                'supabase_response': supabase_response.data
            })
        }
        
    except Exception as e:
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }