delete_user.py

import os
import pickle
import argparse
import shutil
import subprocess

# 경로 설정
DATA_DIR = r'D:\DATA'
train_emb_file = 'train_embeddings.pkl'
val_emb_file = 'val_embeddings.pkl'
label_to_id_file = 'label_to_id.pkl'

def delete_user(pid, remove_folder=False):
    # 파일 로딩
    try:
        with open(train_emb_file, 'rb') as f: train_emb = pickle.load(f)
        with open(val_emb_file, 'rb') as f: val_emb = pickle.load(f)
        with open(label_to_id_file, 'rb') as f: label2id = pickle.load(f)
    except FileNotFoundError:
        print("[오류] 임베딩 또는 라벨 파일이 존재하지 않습니다.")
        return

    if pid not in label2id:
        print(f"[오류] 사용자 '{pid}' 는 등록되어 있지 않습니다.")
        return

    print(f"[삭제] 사용자 '{pid}' 데이터 제거 중...")

    # 삭제
    label2id.pop(pid, None)
    train_emb.pop(pid, None)
    val_emb.pop(pid, None)

    # 저장
    with open(train_emb_file, 'wb') as f: pickle.dump(train_emb, f)
    with open(val_emb_file, 'wb') as f: pickle.dump(val_emb, f)
    with open(label_to_id_file, 'wb') as f: pickle.dump(label2id, f)

    # 폴더 삭제 (옵션)
    user_dir = os.path.join(DATA_DIR, pid)
    if remove_folder and os.path.exists(user_dir):
        shutil.rmtree(user_dir)
        print(f"[삭제] 사용자 폴더도 제거됨: {user_dir}")

    # 자동 재학습
    print("[재학습] 사용자 제거 반영 중...")
    try:
        subprocess.run(['python', 'train_arcface.py'], check=True)
        print("✅ 분류기 재학습 완료.")
    except subprocess.CalledProcessError:
        print("⚠️ 분류기 재학습 중 오류 발생!")

    print(f"✅ 사용자 '{pid}' 데이터 삭제 완료.")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--pid', type=str, required=True, help='삭제할 사용자 ID')
    parser.add_argument('--remove_folder', action='store_true', help='해당 폴더까지 삭제할 경우 지정')
    args = parser.parse_args()

    delete_user(args.pid, args.remove_folder)