from PIL import Image
from elasticsearch import Elasticsearch
from flask import Flask, request, jsonify, render_template
from img2vec_pytorch import Img2Vec
import io
import os

DEFAULT_INDEX = "img-test"
app = Flask(__name__)
es = Elasticsearch(os.environ.get("ES_SERVER") or "http://localhost:9200")


def rgba2rgb(image_file):
    # Open the image file
    img = Image.open(image_file)

    # Check if the image is in RGBA mode
    if img.mode == "RGBA":
        # Convert the image to RGB mode
        img = img.convert("RGB")

        # Create a BytesIO object and save the image to it
        image_io = io.BytesIO()
        img.save(image_io, format="JPEG")

        # Seek to the beginning of the BytesIO object
        image_io.seek(0)

        return image_io

    return image_file


def vectorize(input):
    img2vec = Img2Vec()
    try:
        img = Image.open(input)
        vec = img2vec.get_vec(img, tensor=True)
        vec_np = vec.cpu().numpy().flatten().tolist()
        return vec_np
    except Exception as e:
        print(f"Error processing image: {e}")
        return None


def init_indicies(index: str):
    if es.indices.exists(index):
        return
    # 初始化 kNN 索引
    print(f"Initializing {index}")
    es.indices.create(
        index,
        body={
            "settings": {"index.knn": True},
            "mappings": {
                "properties": {
                    "my_vec": {
                        "type": "knn_dense_float_vector",
                        "knn": {
                            "dims": 512,
                            "model": "lsh",
                            "similarity": "cosine",
                            "L": 99,
                            "k": 1,
                        },
                    }
                }
            },
        },
    )
    img_dir = "static/img"
    for title in os.listdir(img_dir):
        print(f"Indexing {title}")
        my_vec = vectorize(os.path.join(img_dir, title))
        body = {"title": title, "my_vec": my_vec}
        es.index(index=index, body=body)


@app.route("/search", methods=["POST"])
def search_service():
    # 获取表单数据
    index_name = request.form.get("index_name") or DEFAULT_INDEX  # 索引名

    # 获取上传的图片文件
    image_file = request.files.get("image")

    if not index_name or not image_file:
        return jsonify({"error": "Both index_name and image are required."}), 400

    # 处理图片
    image0 = rgba2rgb(image_file)
    vector_arr = vectorize(image0)

    if vector_arr is None:
        return jsonify({"error": "Error processing image."}), 400

    query_body = {
        "size": 50,
        "_source": "title",
        "query": {
            "bool": {
                "must": [
                    {
                        "knn_nearest_neighbors": {
                            "field": "my_vec",
                            "vec": {"values": vector_arr},
                            "model": "lsh",
                            "similarity": "cosine",
                            "candidates": 50,
                        }
                    }
                ]
            }
        },
    }

    if not index_name or not vector_arr:
        return jsonify({"error": "Both index_name and query are required."}), 400

    # 执行搜索
    response = es.search(index=index_name, body=query_body)

    # 使用模板显示搜索结果
    results = response["hits"]["hits"]
    print([r["_source"]["title"] for r in results], len(results))
    return render_template("search.html", results=results)


@app.route("/", methods=["GET"])
def home():
    return render_template("home.html")


if __name__ == "__main__":
    init_indicies(DEFAULT_INDEX)
    app.run(port=5000, debug=True)

 <style>
      #searchForm {
        background-image: url("/static/background/bluewhite.jpg");
        background-repeat: no-repeat;
        background-size: cover;
      }
      body {
        font-family: Arial, sans-serif;
        background-color: #f9f8f8;
        margin: 0;
        padding: 20px;
      }
      .searchImage {
        max-width: 600px;
        max-height: 500px;
      }

      .container {
        max-width: 100%;
        margin: 0 auto;
        /* background: linear-gradient(to right, #8aa0ee, #3838ee); */
        background: #fff;
        padding: 20px;
        border-radius: 5px;
        box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
      }

      h1 {
        text-align: center;
        margin-bottom: 20px;
        color: #000;
      }

      h2 {
        margin-bottom: 10px;
      }

      form {
        max-width: 400px;
        margin: 0 auto;
        background-color: #fff;
        padding: 20px;
        border-radius: 5px;
        box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
      }

      label {
        display: block;
        font-weight: bold;
        margin-bottom: 5px;
      }

      input[type="text"],
      input[type="file"] {
        width: 100%;
        padding: 8px;
        margin-bottom: 10px;
        border: 1px solid #ccc;
        border-radius: 4px;
        box-sizing: border-box;
      }

      input[type="submit"] {
        width: 100%;
        padding: 10px;
        background-color: #4caf50;
        color: #fff;
        border: none;
        border-radius: 4px;
        cursor: pointer;
        transition: background-color 0.3s ease;
      }

      input[type="submit"]:hover {
        background-color: #45a049;
      }

      pre {
        background-color: #f5f5f5;
        padding: 10px;
        border-radius: 5px;
        white-space: pre-wrap;
      }

      #result {
        text-align: center;
        margin-top: 20px;
      }
    </style>
    <style>
      body {
        font-family: Arial, sans-serif;
        /* background-color: #f2f2f2; */
        margin: 0;
        padding: 20px;
      }

      .container {
        max-width: 100%;
        margin: 0 auto;
        background-image: url("/static/background/daziran.png");
      }

      .waterfall-container {
        display: flex;
        flex-wrap: wrap;
      }

      .waterfall-item {
        display: inline-block;
        border-radius: 5px;
        box-shadow: none;
        margin-bottom: 20px;
        text-align: center;
        width: 25%; /* 每行显示 4 个搜索结果，可以根据需要进行调整 */
        padding: 10px;
        box-sizing: border-box;
      }

      .waterfall-item img {
        max-width: 100%;
        max-height: 250px;
        border-radius: 5px;
      }

      .waterfall-item p {
        margin-top: 10px;
        font-size: 25px;
        font-weight: bold;
        color: black;
        background-color: white;
      }
    </style>

 <script src="/static/jquery-3.5.1.min.js"></script>
    <script src="/static/m.js"></script>
    <script>
      $(document).ready(function () {
        $("#searchForm").on("submit", function (event) {
          event.preventDefault();
          var indexName = $("#indexName").val();
          // 检查索引名是否为空或未定义
          if (!indexName) {
            indexName = "img-test"; // 设置默认值为 "默认索引名"
          }
          var formData = new FormData();
          formData.append("index_name", indexName);
          formData.append("image", $("#image")[0].files[0]);

          $.ajax({
            url: "/search",
            method: "POST",
            processData: false, // Important!
            contentType: false, // Important!
            data: formData,
            success: function (data) {
              // Clear the previous results
              $("#result").empty();

              // Parse the returned HTML and extract image and paragraph elements
              var parsedData = $.parseHTML(data);

              // Group image and paragraph elements
              var imageAndParagraphPairs = [];
              var currentPair = [];
              $(parsedData).each(function () {
                if ($(this).is("img.searchImage")) {
                  if (currentPair.length === 1) {
                    currentPair.push(this);
                    imageAndParagraphPairs.push(currentPair);
                    currentPair = [];
                  } else {
                    currentPair.push(this);
                  }
                } else if ($(this).is("p")) {
                  if (currentPair.length === 0) {
                    currentPair.push(this);
                  } else {
                    currentPair.push(this);
                    imageAndParagraphPairs.push(currentPair);
                    currentPair = [];
                  }
                } else if ($(this).is("h1")) {
                  // Add the <h1> element back to the results

                  $("#resultTitle").html($(this));
                }
              });

              // Create and append the waterfall items
              $.each(imageAndParagraphPairs, function (index, pair) {
                var $item = $("<div>").addClass("waterfall-item");
                $.each(pair, function (i, element) {
                  $item.append(element);
                });
                $("#result").append($item);
              });
            },
          });
        });
      });
    </script>