semantic_search/amazon_rekognition.py

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