import cv2
import matplotlib.pyplot as plt
import gradio as gr
import os


def find_descriptors(img1, img2):
    if img1 is None or img2 is None:
        return None, None, None, None, None, None

    img1 = cv2.cvtColor(img1, cv2.IMREAD_GRAYSCALE)
    img2 = cv2.cvtColor(img2, cv2.IMREAD_GRAYSCALE)
    sift = cv2.SIFT_create()
    kp1, des1 = sift.detectAndCompute(img1, None)
    kp2, des2 = sift.detectAndCompute(img2, None)
    bf = cv2.BFMatcher()
    matches = bf.knnMatch(des1, des2, k=2)
    good = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good.append([m])
    img3 = cv2.drawMatchesKnn(img1, kp1, img2, kp2, good, None, flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
    plt.imshow(img3), plt.show()
    score = (100 * len(good) / min(len(kp1), len(kp2)))
    if score >= 10:
        result = "Totally YES"
    elif score >= 3:
        result = "YES"
    elif score >= 1:
        result = "MAYBE"
    else:
        result = "NO"

    descriptors = len(good)
    keypoints1 = len(kp1)
    keypoints2 = len(kp2)

    return img3, descriptors, keypoints1, keypoints2, result, score


sift_result = gr.Interface(
    fn=find_descriptors,
    inputs=["image", "image"],
    outputs=["image", gr.Textbox(label="Descriptors"),
             gr.Textbox(label="First image keypoints"),
             gr.Textbox(label="Second image keypoints"),
             gr.Textbox(label="Are these places equals?"),
             gr.Textbox(label="Score")],
    examples=[[
        os.path.join(os.path.dirname(__file__), "images/1.jpg"),
        os.path.join(os.path.dirname(__file__), "images/2.jpg"),
    ], [os.path.join(os.path.dirname(__file__), "images/3.png"),
        os.path.join(os.path.dirname(__file__), "images/4.png")],
        [os.path.join(os.path.dirname(__file__), "images/5.jpg"),
         os.path.join(os.path.dirname(__file__), "images/6.jpg")],
        [os.path.join(os.path.dirname(__file__), "images/7.jpg"),
         os.path.join(os.path.dirname(__file__), "images/8.jpg")]],
    title="SIFT Image Compare",
    batch=False,
    description="A model that uses SIFT method for image compare.",
    live=True)

sift_result.launch()