iframe-proxy | Sunbelt Computer Software

History

183 lines (161 loc) · 6.96 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

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)})

}

Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Uh oh!

Sunbelt Computer Software

PL/B Language Development and Support

Uh oh!

FilesExpand file tree

lambda_function.py

Latest commit

History

lambda_function.py

File metadata and controls