ssim.py

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

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)

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 compare_frames(frame1, frame2):
    """Compare two frames using structural similarity index (SSIM)."""
    frame1 = frame1.convert('L')
    frame2 = frame2.convert('L')
    if frame1.size != frame2.size:
        frame2 = frame2.resize(frame1.size)
    arr1 = np.array(frame1)
    arr2 = np.array(frame2)
    similarity = ssim(arr1, arr2, data_range=255)
    return similarity

def lambda_handler(event, context):
    try:
        # Check if a job_id is provided; if not, generate one.
        job_id = event.get('job_id')
        if not job_id:
            job_id = str(uuid.uuid4())
            # Insert a new record in the ai_jobs table with the generated job_id.
            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', 'png')
        
        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 shortcode (single download per job)
        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 separately
        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 based on the monitored frames order.
            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 records 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)
            })
        }