python_clip_image_search/server.py

import configparser
from flask import Flask, json, request
import logging
import os
import search_engine_model as sem

class SearchEngineApp():

    def __init__(self):
        self.logger = logging.getLogger(__name__)
        logging.basicConfig(level=logging.INFO)
        self.config_manager_obj = configparser.ConfigParser()
        self.config_manager_obj.read('./config.cfg')

        image_root_dir = self.config_manager_obj['SEARCH_ENGINE']['image_root_dir']
        csv_file_path = self.config_manager_obj['SEARCH_ENGINE']['csv_file_path']
        self.model_manager_obj = sem.SearchEngineModel(image_root_dir, csv_file_path)

    # PRECONDITION: We assume that only text_prompt or image_prompt
    # will be different from None.
    def search_images(self, text_prompt, image_prompt, prompt_type, nofimages_to_show):
        search_result = None
        if ((prompt_type=='Text') and (text_prompt!=None) and (len(text_prompt.strip())>0)):
            search_result = self.model_manager_obj.search_image_by_text_prompt(text_prompt, 
                                                                               nofimages_to_show)

        if ((prompt_type=='Image') and (image_prompt!=None)):
            search_result = self.model_manager_obj.search_image_by_image_prompt(image_prompt,
                                                                                nofimages_to_show)

        return search_result

# Flask application code.

api = Flask(__name__)
# I use this as a global variable, avoiding a model load each time a request is made.
search_engine_app_manager = SearchEngineApp()

@api.route('/images', methods=['GET'])
def get_images():
    text_prompt = request.args.get('search_input')        
    search_result = search_engine_app_manager.search_images(text_prompt, None, 'Text', 20)

    return json.dumps({'images': search_result})

if __name__ == '__main__':
    api.run(host='0.0.0.0', port=5000, debug=True)