Uh oh!
There was an error while loading. Please reload this page.
- Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathlambda_function.py
More file actions
Latest commit
183 lines (161 loc) · 6.96 KB
/
Copy pathlambda_function.py
File metadata and controls
183 lines (161 loc) · 6.96 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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
importos
importjson
importtime
importuuid
importboto3
fromPILimportImage
importnumpyasnp
fromioimportBytesIO
fromconcurrent.futuresimportThreadPoolExecutor, as_completed
fromsupabaseimportcreate_client, Client
fromdotenvimportload_dotenv
importtorch
importclip
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"iftorch.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)
defdownload_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)
returnImage.open(BytesIO(response['Body'].read()))
excepts3_client.exceptions.NoSuchKey:
returnNone
defget_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
withThreadPoolExecutor() asexecutor:
future_to_frame= {
executor.submit(download_frame, shortcode, i, platform, extension): i
foriinrange(max_frames)
}
forfutureinas_completed(future_to_frame):
frame_number=future_to_frame[future]
frame=future.result()
ifframe:
frames.append((frame_number, frame))
# Sort frames by frame number to maintain order
frames.sort(key=lambdax: x[0])
return [framefor_, frameinframes]
defextract_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)
withtorch.no_grad():
features=model.encode_image(image_input)
returnfeatures.cpu().numpy().squeeze()
defcosine_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)
ifnorm_a==0ornorm_b==0:
return0.0
returndot_product/ (norm_a*norm_b)
defcompare_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)
returnsimilarity
deflambda_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')
ifnotjob_id:
job_id=str(uuid.uuid4())
job_insert_response=supabase.table("ai_jobs").insert({"job_id": job_id}).execute()
ifjob_insert_response.error:
raiseException(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')
ifnotmonitored_shortcodeornotwatched_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)
ifnotmonitored_frames:
return {
'statusCode': 404,
'body': json.dumps({'error': f'No frames found for monitored shortcode: {monitored_shortcode}'})
}
records_to_insert= []
# Process each watched video
forwatched_shortcodeinwatched_shortcodes:
start_time_video=time.time()
watched_frames=get_all_frames(watched_shortcode, platform, extension)
ifnotwatched_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
fori, (monitored_frame, watched_frame) inenumerate(zip(monitored_frames, watched_frames)):
similarity=compare_frames(monitored_frame, watched_frame)
frame_comparisons.append({'frame_number': i, 'similarity': float(similarity)})
similarities= [comp['similarity'] forcompinframe_comparisons]
avg_similarity=float(np.mean(similarities)) ifsimilaritieselseNone
max_similarity=float(max(similarities)) ifsimilaritieselseNone
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
})
}
exceptExceptionase:
return {
'statusCode': 500,
'body': json.dumps({'error': str(e)})
}