import cv2
import numpy as np
import glob
import os

def rotate_image(image, angle):
    """画像を指定された角度で回転させる"""
    center = (image.shape[1] // 2, image.shape[0] // 2)
    rot_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated_image = cv2.warpAffine(image, rot_matrix, (image.shape[1], image.shape[0]))
    return rotated_image

def scale_image(image, scale):
    """画像を指定されたスケールで拡大・縮小する"""
    width = int(image.shape[1] * scale)
    height = int(image.shape[0] * scale)
    scaled_image = cv2.resize(image, (width, height))
    return scaled_image

def zncc(template, image):
    """テンプレートと画像のZNCCを計算"""
    template = (template - np.mean(template)) / np.std(template)
    image = (image - np.mean(image)) / np.std(image)
    result = cv2.matchTemplate(image.astype(np.float32), template.astype(np.float32), cv2.TM_CCORR_NORMED)
    return result

# テンプレート画像を読み込む
template_files = glob.glob('template_img/template*.JPG')
templates = [cv2.imread(template_file, cv2.IMREAD_GRAYSCALE) for template_file in template_files]

# テンプレート画像が読み込めたか確認
for template_file, template in zip(template_files, templates):
    if template is None:
        raise FileNotFoundError(f"Template image {template_file} could not be read.")

# 畑の画像の読み込み
image = cv2.imread('input_img/field.JPG', cv2.IMREAD_GRAYSCALE)
image_color = cv2.imread('input_img/field.JPG')
if image is None:
    raise FileNotFoundError("Field image could not be read.")

# 回転角度のリスト
angles = range(-10, 10, 1)  # -10度から10度まで1度刻みで回転
# スケールのリスト
scales = [0.1, 0.2, 0.3, 0.4]  # テンプレートを10%、20%、30%、40%にスケール

# しきい値の設定
threshold = 0.50

# 各テンプレート、回転角度、スケールに対してZNCCを計算
detections = []

for template_index, template in enumerate(templates):
    for scale in scales:
        scaled_template = scale_image(template, scale)
        h, w = scaled_template.shape
        for angle in angles:
            rotated_template = rotate_image(scaled_template, angle)
            result = zncc(rotated_template, image)
            
            # しきい値を超える場所を見つける
            loc = np.where(result >= threshold)
            for pt in zip(*loc[::-1]):  # Y座標、X座標を入れ替える
                if not any(abs(pt[0] - d[0][0]) <= 2 and abs(pt[1] - d[0][1]) <= 2 for d in detections):
                    detections.append((pt, (w, h)))  # 座標とテンプレートサイズを保存
        print(len(detections))

if detections:
    print(f"収穫可能なきゅうりが{len(detections)}個見つかりました")
    for (top_left, (w, h)) in detections:
        bottom_right = (top_left[0] + w, top_left[1] + h)
        cv2.rectangle(image_color, top_left, bottom_right, (0, 0, 255), 5)  # 見つけた領域を赤の矩形で囲む
    
    # outputフォルダを作成（存在しない場合）
    output_dir = 'output'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 画像を保存
    output_path = os.path.join(output_dir, 'detected_cucumbers.jpg')
    cv2.imwrite(output_path, image_color)
    print(f"検出された画像を {output_path} に保存しました")
else:
    print("収穫可能なきゅうりは見つかりませんでした")