OpenSearch-AI / semantic_search /amazon_rekognition.py
prasadnu's picture
rerank model
eb03410
raw
history blame
3.69 kB
import json
import os
import sys
import boto3
from botocore.config import Config
import getpass
import streamlit as st
from opensearchpy import OpenSearch, RequestsHttpConnection
rekog_client = boto3.client('rekognition',aws_access_key_id=st.secrets['user_access_key'],
aws_secret_access_key=st.secrets['user_secret_key'],region_name='us-east-1')
def extract_image_metadata(img):
res = rekog_client.detect_labels(
Features= [ "GENERAL_LABELS","IMAGE_PROPERTIES" ],
Image = {
"Bytes":img
},
MaxLabels = 10,
MinConfidence = 80.0,
Settings = {
"ImageProperties": {
"MaxDominantColors": 5
}
}
)
objects_category_color = {}
objects_category_color_masked = {}
def add_span(x,type):
if(type == 'obj'):
return "<span style='color:#e28743;font-weight:bold'>"+x+"</span>"
if(type == 'cat'):
return "<span style='color:green;font-style:italic'>"+x+"</span>"
if(type == 'color'):
return "<span style='color:"+x+";font-weight:bold'>"+x+"</span>"
for label in res['Labels']:
objects_category_color_masked[add_span(label['Name'],'obj')]={'categories':[],'color':""}
objects_category_color[label['Name']] = {'categories':[],'color':""}
if(len(label['Categories'])!=0):
for category in label['Categories']:
objects_category_color_masked[add_span(label['Name'],'obj')]['categories'].append(add_span(category['Name'].lower(),'cat'))
objects_category_color[label['Name']]['categories'].append(category['Name'].lower())
if(len(label['Instances'])!=0):
for instance in label['Instances']:
if(len(instance['DominantColors'])!=0):
objects_category_color[label['Name']]['color'] = instance['DominantColors'][0]['SimplifiedColor']
objects_category_color_masked["<span style='color:#e28743;font-weight:bold'>"+label['Name']+"</span>"]['color'] = add_span(instance['DominantColors'][0]['SimplifiedColor'],'color')
st.session_state.input_rekog_directoutput = objects_category_color_masked
objects = []
categories = []
colors = []
for key in objects_category_color.keys():
if(key.lower() not in objects):
objects.append(key.lower())
categories.append(" ".join(set(objects_category_color[key]['categories'])))
if(objects_category_color[key]['color']!=''):
colors.append(objects_category_color[key]['color'].lower())
objects = " ".join(set(objects))
categories = " ".join(set(categories))
colors = " ".join(set(colors))
return colors+ " " + objects + " " + categories
def call(a,b):
if(st.session_state.input_is_rewrite_query == 'enabled' and st.session_state.input_rewritten_query!=""):
st.session_state.input_rewritten_query['query']['bool']['should'].append( {
"simple_query_string": {
"query": a + " " + b,
"fields":['description','rekog_all^3']
}
})
rekog_query = st.session_state.input_rewritten_query
else:
rekog_query = { "query":{
"simple_query_string": {
"query": a + " " + b,
"fields":['description','rekog_all^3']
}
}
}
st.session_state.input_rewritten_query = rekog_query