import cv2
import numpy as np

import khandy
from insectid import InsectDetector, InsectIdentifier

detector = InsectDetector()
identifier = InsectIdentifier()


def query(img_file):
    insect_image = khandy.imread(img_file)  # Read image
    if insect_image is None:
        print('Failed to read image.')
    if max(insect_image.shape[:2]) > 1280:  # Resize image if it's too large
        insect_image = khandy.resize_image_long(insect_image, 1280)

    image_for_draw = insect_image.copy()  # Copy image for drawing
    image_height, image_width = insect_image.shape[:2]

    boxes, confs, classes = detector.detect(insect_image)
    for box, conf, class_ind in zip(boxes, confs, classes):  # Loop through all detected objects
        box = box.astype(np.int32)
        box_width = box[2] - box[0] + 1
        box_height = box[3] - box[1] + 1
        if box_width < 30 or box_height < 30:
            continue

        cropped = khandy.crop(insect_image, box[0], box[1], box[2], box[3])
        results = identifier.identify(cropped)

        print(results[0])  # 打印第一个结果，这里可以根据需要进行处理
        prob = results[0]['probability']

        if prob < 0.10:
            text = 'Unknown'
        else:
            text = '{}: {:.3f}'.format(results[0]['chinese_name'], results[0]['probability'])

        # 获取文本位置
        position = [box[0] + 2, box[1] - 20]
        position[0] = min(max(position[0], 0), image_width)
        position[1] = min(max(position[1], 0), image_height)

        # 画框
        cv2.rectangle(image_for_draw, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
        image_for_draw = khandy.draw_text(image_for_draw, text, position, font_size=15)
    return image_for_draw


if __name__ == '__main__':
    image_file = 'images/七星瓢虫.jpg'
    image = query(image_file)
    # 保存图片到文件
    cv2.imwrite('images/query_result.jpg', image)