validate_inference.py

import os
import json
import time
import argparse
import importlib
import numpy as np
from collections import Counter

from stereo_undistort_3.undistort_single import undistort_stereo_images
from object_detection.model import detect_objects
from layer_estimation.model import layer_estimation
from layer_estimation.visualization import visualize_layer_result

def logging_result(args, base_name, gt_floor, pred_results):
    """
    예측 값과 정답 값을 base_name 별로 json 형식의 파일에 저장
    
    Args:
        base_name : 입력된 파일 이름
        gt_floor : 정답 층수 배열
        pred_results : 예측 결과 객체
    """
    data = {}

    pred_floor = [d['layer'] for d in pred_results]
    estimate_height = [d['depth_m'] for d in pred_results]  # 추론 높이
    gt_height = [0.15 * f for f in gt_floor]                # 실제 높이

    result_json_filepath = f"{args.output_dir}/{args.run_name}/result.json"

    if os.path.exists(result_json_filepath):
        with open(result_json_filepath, 'r', encoding='utf-8') as f:
            file_content = f.read().strip()
            data = json.loads(file_content) if file_content else {}
    
    
    new_entry = {
        "gt_floor": gt_floor,
        "pred_floor": pred_floor,
        "estimate_height": estimate_height,
        "gt_height": gt_height,
        "diff_height": [abs(gt - pred) for gt, pred in zip(gt_height, estimate_height)]
    }

    data[base_name] = new_entry 
    
    with open(result_json_filepath, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=4)

def evaluate_per_image(gt_floor, pred_floor):
    """
    이미지 단위로 평가
    """
    gt_floor = np.array(gt_floor)
    pred_floor = np.array(pred_floor)

    obj_accuracy = np.mean(pred_floor == gt_floor)   # 객체 단위 정확도
    all_correct = np.all(pred_floor == gt_floor)     # 모든 객체가 맞았는가
    mae = np.mean(np.abs(pred_floor - gt_floor))     # 평균 절대 오차
    rmse = np.sqrt(np.mean((pred_floor - gt_floor) ** 2))  # RMSE

    gt_count = Counter(gt_floor)
    pred_count = Counter(pred_floor)

    """
    TP: 예측 박스 중 실제 정답 박스와 '같은 층으로' 매칭된 경우
    FP: 실제에는 없는데 잘못 탐지된 박스
    FN: 실제에는 있는데 탐지하지 못한 박스
    """
    tp = sum(min(gt_count[layer], pred_count[layer]) for layer in set(gt_count.keys()) | set(pred_count.keys()))
    fp = len(pred_floor) - tp
    fn = len(gt_floor) - tp

    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0


    return {
        "obj_accuracy": obj_accuracy,
        "all_correct": all_correct,
        "mae": mae,
        "rmse": rmse,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "tp": tp,
        "fp": fp,
        "fn": fn
    }

def evaluate_from_logged_results(args):
    """
    전체 test set 성능을 계산
    """
    result_json_filepath = f"{args.output_dir}/{args.run_name}/result.json"
    
    with open(result_json_filepath, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 기존 데이터에서 이미지별 결과만 추출 (최종 통계가 이미 있는 경우 제외)
    image_entries = {
        k: v for k, v in data.items()
        if isinstance(v, dict) and "gt_floor" in v and "pred_floor" in v
    }

    all_obj_accuracies, all_mae, all_rmse = [], [], []
    all_image_correct, all_precisions, all_recalls, all_f1 = [], [], [], []
    total_tp, total_fp, total_fn = 0, 0, 0

    for base_name, result in image_entries.items():
        gt_floor = result["gt_floor"]
        pred_floor = result["pred_floor"]

        # ground truth가 비어 있으면 스킵
        if len(gt_floor) == 0:
            print(f"[SKIP] {base_name}: gt_floor 비어 있음 (pred={len(pred_floor)})")
            continue

        if len(gt_floor) != len(pred_floor):
            print(f"[SKIP] {base_name}: gt_floor와 pred_floor 개수 불일치 (gt={len(gt_floor)}, pred={len(pred_floor)})")
            continue
        
        metrics = evaluate_per_image(gt_floor, pred_floor)
        all_obj_accuracies.append(metrics["obj_accuracy"])
        all_mae.append(metrics["mae"])
        all_rmse.append(metrics["rmse"])
        all_image_correct.append(metrics["all_correct"])
        all_precisions.append(metrics["precision"])
        all_recalls.append(metrics["recall"])
        all_f1.append(metrics["f1"])

        total_tp += metrics["tp"]
        total_fp += metrics["fp"]
        total_fn += metrics["fn"]

        print(f"[{base_name}] 객체정확도: {metrics['obj_accuracy']:.2f}, F1: {metrics['f1']:.2f}, "
              f"TP={metrics['tp']}, FP={metrics['fp']}, FN={metrics['fn']}")

    mean_obj_acc = np.nanmean(all_obj_accuracies) * 100
    mean_mae = np.nanmean(all_mae)
    mean_rmse = np.nanmean(all_rmse)
    image_accuracy = (np.sum(all_image_correct) / len(all_image_correct)) * 100
    mean_precision = np.mean(all_precisions)
    mean_recall = np.mean(all_recalls)
    mean_f1 = np.mean(all_f1)

    # --- 전체 TP/FP/FN 기반 최종 F1 ---
    total_precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0
    total_recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0
    total_f1 = (2 * total_precision * total_recall) / (total_precision + total_recall) if (total_precision + total_recall) > 0 else 0

    print("\n=== 전체 테스트 세트 결과 ===")
    print(f"객체 단위 평균 정확도: {mean_obj_acc:.2f}%")
    print(f"이미지 단위 전체 일치 비율: {image_accuracy:.2f}%")
    print(f"평균 MAE: {mean_mae:.3f}")
    print(f"평균 RMSE: {mean_rmse:.3f}")
    print(f"평균 Precision: {mean_precision:.3f}")
    print(f"평균 Recall: {mean_recall:.3f}")
    print(f"평균 F1 Score (이미지 단위 평균): {mean_f1:.3f}")
    print(f"최종 Detection F1 Score (전체 TP/FP/FN 기준): {total_f1:.3f}")

    data["summary"] = {
        "run_name": args.run_name,
        "undistortion_alpha": args.alpha,
        "use_undistortion": args.use_undistortion,
        "use_pre_bbox": args.use_pre_bbox,
        "depth_model": args.depth_model,
        "use_pre_disp": args.use_pre_disp,
        "use_flatten": args.use_flatten,
        "layer_inference_method": args.method,
        "obj_accuracy(%)": mean_obj_acc,
        "image_accuracy(%)": image_accuracy,
        "mae": mean_mae,
        "rmse": mean_rmse,
        "mean_precision": mean_precision,
        "mean_recall": mean_recall,
        "mean_f1": mean_f1,
        "final_f1": total_f1,
        "num_images": len(image_entries),
        "total_tp": total_tp,
        "total_fp": total_fp,
        "total_fn": total_fn
    }

    with open(result_json_filepath, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=4)

    print(f"\n📁 결과가 '{result_json_filepath}' 파일에 저장되었습니다.")

    return {
        "run_name": args.run_name,
        "undistortion_alpha": args.alpha,
        "use_undistortion": args.use_undistortion,
        "use_pre_bbox": args.use_pre_bbox,
        "depth_model": args.depth_model,
        "use_pre_disp": args.use_pre_disp,
        "use_flatten": args.use_flatten,
        "layer_inference_method": args.method,
        "obj_accuracy": mean_obj_acc,
        "image_accuracy": image_accuracy,
        "mae": mean_mae,
        "rmse": mean_rmse,
        "precision": mean_precision,
        "recall": mean_recall,
        "mean_f1": mean_f1,
        "final_f1": total_f1
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--use_undistortion', default=False, action="store_true")
    parser.add_argument('--use_pre_bbox', default=False, action="store_true")
    parser.add_argument('--depth_model', default='defom')
    parser.add_argument('--use_pre_disp', default=False, action="store_true")
    parser.add_argument('--use_flatten', default=False, action="store_true")
    parser.add_argument('--method', default='center', help="inference method")
    parser.add_argument('--verbose', default=False, action="store_true")
    parser.add_argument('--alpha', default=1.0, type=float)
    parser.add_argument('--target_video_name', default=None)
    parser.add_argument('--calib_dir_num', type=int, default=3)
    args = parser.parse_args()
    
    base_dir = "/abr/coss32/workspace/dataset/layer_dataset"  # FIXME
    image_dataset_dir = f"{base_dir}/images"
    
    # 실험 폴더 및 이름 지정
    args.output_dir = "/abr/coss32/workspace/result"
    args.run_name = time.strftime("%Y%m%d-%H%M%S")

    try:
        result_json_filepath = f"{args.output_dir}/{args.run_name}/result.json"
        os.makedirs(os.path.dirname(result_json_filepath), exist_ok=True)
    except OSError as e:
        print(f"결과 폴더 생성 오류 발생: {e}")

    if args.verbose:
        print(f"Start {args.run_name}")
        print(f"use_undistortion: {args.use_undistortion} with alpha={args.alpha}")
        print(f"use_pre_bbox: {args.use_pre_bbox}")
        print(f"depth_model: {args.depth_model}")
        print(f"use_pre_disp: {args.use_pre_disp}")
        print(f"use_flatten: {args.use_flatten}")
        print(f"method: {args.method}")

    video_names = [d for d in os.listdir(image_dataset_dir) if os.path.isdir(os.path.join(image_dataset_dir, d))]

    for video_name in video_names:
        # FIXME video_name 제한
        if video_name not in ["2x2x3_random_1", "2x3x3_one", "2x3x3_one_1", "2x3x3_random", "3x3x3_one", "3x3x3_random"]:
            continue

        if args.target_video_name is not None and video_name != args.target_video_name:
            continue

        image_dir = f"{image_dataset_dir}/{video_name}"
        bbox_npy_dir = f"{base_dir}/bbox/{video_name}"
        disp_map_dir = f"{base_dir}/disparity_map_improve/{video_name}/{args.depth_model}/{args.depth_model}_2"
        gt_txt_dir = f"{base_dir}/layers/{video_name}"

        print(f"---------- video: {video_name}, len: {len(os.listdir(image_dir))} ----------")

        # image_list = os.listdir(image_dir)
        # random_n = 10 # 추출할 이미지 수 
        # random_image_list = random.sample(image_list, random_n)

        for image_filename in os.listdir(image_dir):
            target_left_image_path = os.path.join(image_dir, image_filename)
            target_right_image_path = os.path.join(f"/abr/coss32/workspace/dataset/extracted_images/right/{video_name}/{image_filename}")
            bbox_npy_dir = f"{base_dir}/bbox/{video_name}"

            if args.verbose:
                print(f"left image: {target_left_image_path}")
                print(f"right image: {target_right_image_path}")
            
            # target image와 대응되는 파일들을 가져옴
            base_name, _ = os.path.splitext(image_filename)
            if base_name == ".DS_Store":
                continue
            
            target_bbox_path = os.path.join(bbox_npy_dir, f"{base_name}.npy")
            target_disparity_path = os.path.join(disp_map_dir, f"{base_name}.npy")
            target_layer_path = os.path.join(gt_txt_dir, f"{base_name}.txt")

            if not os.path.exists(target_layer_path):
                print(f"[SKIP] GT Layer 파일 없음: {target_layer_path}")
                continue

            # 1단계 : 입력 이미지 불러와서 왜곡 보정
            """
            calib_dir="/abr/coss32/workspace/stereo_undistort/calibration_results"

            calib_dir="/abr/coss32/workspace/stereo_undistort_2/calibration_results"

            calib_dir="/abr/coss32/workspace/stereo_undistort_3/calibration_results",

            calib_dir="/abr/coss32/workspace/stereo_undistort_3/calibration_results_60mm"
            """
            if args.use_undistortion:
                # 어제 저녁에 보낸게 3번으로 돌린겁니다
                _calib_dir="/abr/coss32/workspace/stereo_undistort_3/calibration_results"
                if args.calib_dir_num == 1:
                    print("calib 1")
                    _calib_dir="/abr/coss32/workspace/stereo_undistort/calibration_results"
                elif args.calib_dir_num == 2:
                    print("calib 2")
                    _calib_dir="/abr/coss32/workspace/stereo_undistort_2/calibration_results"
                elif args.calib_dir_num == 3:
                    print("calib 3")
                    _calib_dir="/abr/coss32/workspace/stereo_undistort_3/calibration_results"
                elif args.calib_dir_num == 4:
                    print("calib 4")
                    _calib_dir="/abr/coss32/workspace/stereo_undistort_3/calibration_results_60mm"


                args.left_imgs, args.right_imgs = undistort_stereo_images(
                    target_left_image_path,
                    target_right_image_path,
                    calib_dir=_calib_dir,
                    alpha=args.alpha,
                )
            else:
                # 보정 없이 원본 이미지 사용
                args.left_imgs, args.right_imgs = target_left_image_path, target_right_image_path

            # 2단계 : 객체 인식
            if args.use_pre_bbox:
                try:
                    bbox_npy = np.load(target_bbox_path)
                except FileNotFoundError:
                    print(f"[SKIP] BBox 파일 없음: {target_bbox_path}")
                    continue
                except Exception as e:
                    print(f"[경고] BBox 로드 실패 ({target_bbox_path}) - {e}")
                    bbox_npy = None
            else:
                bbox_npy, _ = detect_objects(image=args.left_imgs, output_dir=args.output_dir, run_name=args.run_name)
            
            # 3단계 : 깊이 추론
            if args.use_pre_disp:
                try:
                    disp_pr = np.load(target_disparity_path)
                except FileNotFoundError:
                    print(f"[SKIP] Disparity 파일 없음: {target_disparity_path}")
                    continue
                except Exception as e:
                    print(f"[경고] Disparity map 로드 실패 ({target_disparity_path}) - {e}")
                    disp_pr = None
            else:
                model_paths = {
                    "defom": "depth_estimation.inf_depth_model.DEFOM_Stereo.model",
                    "mocha": "depth_estimation.inf_depth_model.MoCha_Stereo.model",
                    "selective": "depth_estimation.inf_depth_model.Selective_Stereo.model",
                    "sgbm": "depth_estimation.inf_depth_model.SGBM.model",
                    "stereobm": "depth_estimation.inf_depth_model.StereoBM.model",
                    "finetuned": "depth_estimation.fintune_defom.Song_DEFOM_Stereo.model",
                }
                model_key = args.depth_model.lower()
                if model_key not in model_paths:
                    raise ValueError(f"Invalid depth model: {model_key}. Choose from {list(model_paths.keys())}")

                # 동적 import
                model_module = importlib.import_module(model_paths[model_key])
                depth_estimation = getattr(model_module, "depth_estimation")
                disp_pr, _ = depth_estimation(args, output_dir=args.output_dir, run_name=args.run_name)

            # 3.5단계 : bbox 단위 depth 평면화
            if args.use_flatten:
                from depth_flatten import flatten_bbox_depths
                disp_pr = flatten_bbox_depths(disp_pr, bbox_npy)  # 평면화된 depth map 반환

            # 4단계 : 층수 추론
            out_json, _ = layer_estimation(args, disparity_map_npy=disp_pr, bbox_npy=bbox_npy)

            final_img = visualize_layer_result(args, bbox_npy=bbox_npy, layer_json=out_json, output_dir=args.output_dir, run_name=args.run_name)

            gt_floor = []
            with open(target_layer_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            data_dict = json.loads(line)
                            gt_floor.append(data_dict["layer"])
                        except json.JSONDecodeError as e:
                            print(f"경고: JSON 파싱 오류 발생 - 줄: '{line}', 오류: {e}")
                        except KeyError:
                            print(f"경고: 'layer' 키가 없습니다 - 줄: '{line}'")
            
            pred_results, _ = layer_estimation(args, disp_pr, bbox_npy)

            # 파일에 추가
            logging_result(args, base_name, gt_floor, pred_results)
    
    if args.verbose:
        print("Summary..")
    summary = evaluate_from_logged_results(args)