import sensor, image, time, os

# 全局标志位（默认未检测到）
recognition_flag = 0
matched_label = "Unknown"

# 初始化摄像头
sensor.reset()
sensor.set_pixformat(sensor.GRAYSCALE)
sensor.set_framesize(sensor.VGA)
sensor.set_windowing((320, 240))
sensor.set_contrast(3)
sensor.skip_frames(time=3000)

# 特征库路径，修改为目录路径
FEATURE_PATH = "singtown/face/"

# 加载本地特征库
known_features = []
known_labels = []

print(f"Current path: {os.getcwd()}")  # 打印当前工作目录
print(f"Checking if {FEATURE_PATH} exists...")
try:
    os.stat(FEATURE_PATH)  # 检查目录是否存在
    print(f"{FEATURE_PATH} exists.")
except OSError:
    print(f"{FEATURE_PATH} does not exist.")
    while True:
        time.sleep(1000)

try:
    print(f"Attempting to list files in {FEATURE_PATH}")
    files = os.listdir(FEATURE_PATH)
    print(f"Files in {FEATURE_PATH}: {files}")
    for f in files:
        if f.endswith(".orb"):
            file_path = FEATURE_PATH + f
            print(f"Attempting to load {file_path}")
            kpts = image.load_descriptor(file_path)
            known_features.append(kpts)
            known_labels.append(f.split(".")[0])
            print(f"Loaded: {f}")
except OSError as e:
    if e.errno == 22:
        print(f"Path error: {FEATURE_PATH} is invalid. Check path format, SD card and file integrity. Flag: 0")
    else:
        print(f"Load error: {e}. Check the path and permissions. Flag: 0")
    while True:
        time.sleep(1000)

face_cascade = image.HaarCascade("frontalface", stages=25)
MATCH_THRESHOLD = 50  # 匹配点数阈值
EXTEND_PIXELS = 31    # ROI扩展像素

def recognize_face():
    global recognition_flag, matched_label
    while True:
        img = sensor.snapshot()
        recognition_flag = 0
        matched_label = "Unknown"

        faces = img.find_features(face_cascade, threshold=0.5, scale_factor=1.2)
        if not faces:
            continue

        for face in faces:
            x, y, w, h = face
            roi = (
                max(0, x - EXTEND_PIXELS),
                max(0, y - EXTEND_PIXELS),
                min(img.width(), w + 2 * EXTEND_PIXELS),
                min(img.height(), h + 2 * EXTEND_PIXELS)
            )

            kpts = img.find_keypoints(
                roi=roi,
                threshold=10,
                scale_factor=1.1,
                max_keypoints=200,
                corner_detector=image.CORNER_AGAST
            )

            if not kpts:
                img.draw_rectangle(face, color=(255, 0, 0))
                continue

            best_match = -1
            best_count = 0
            for i, known_kpts in enumerate(known_features):
                match_count = image.match_descriptor(kpts, known_kpts, threshold=50)
                if match_count > best_count:
                    best_count = match_count
                    best_match = i

            if best_count >= MATCH_THRESHOLD:
                recognition_flag = 1
                matched_label = known_labels[best_match]
                img.draw_rectangle(face, color=(0, 255, 0), thickness=2)
                img.draw_string(x, y - 15, f"ID:{matched_label}", color=(0, 255, 0))
                print(f"[+] Match: {matched_label} (Flag: {recognition_flag})")
            else:
                img.draw_rectangle(face, color=(255, 0, 0), thickness=2)
                img.draw_string(x, y - 15, "Unknown", color=(255, 0, 0))
                print(f"[-] No match (Flag: {recognition_flag})")

# 主程序入口
if __name__ == "__main__":
    if not known_features:
        print("Error: No feature files found! Flag: 0")
    else:
        recognize_face()