src/evaluate_ne.py

# evaluate_ne.py

import sys
import os
import cv2
import numpy as np
import time
import argparse
from glob import glob

# Thêm thư mục cha vào path để import được src
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from src import utils
from src import ml_patch

def run_ne_inference(lr_img_path, knn_model, feat_vectors, hr_patches_db, scale, patch_size, overlap, k_neighbors, epsilon):
    """
    Hàm tái tạo ảnh NE cho một ảnh LR đầu vào.
    Trả về: Ảnh kênh Y (float 0-255)
    """
    # 1. Trích xuất patches và features từ ảnh LR
    _, target_LR_feats, target_lr_y_means = ml_patch.extract_patches_per_image(
        lr_img_path, scale=scale, patch_size=patch_size, overlap=overlap
    )
    
    predicted_patches = []
    
    # 2. Duyệt qua từng feature patch để tìm láng giềng
    # (Lưu ý: Có thể tối ưu bằng batch query nhưng loop dễ hiểu hơn)
    for i, query_feat in enumerate(target_LR_feats):
        # KNN Search
        distances, indices = ml_patch.query_knn(knn_model, feat_vectors, query_feat, k=k_neighbors)
        
        neighbor_vecs = feat_vectors[indices]
        neighbor_hr = hr_patches_db[indices]
        
        # Tính trọng số LLE
        w = ml_patch.compute_weights(query_feat, neighbor_vecs, epsilon=epsilon)
        
        # Tái tạo patch HR
        pred_patch = ml_patch.map_hr_patches(neighbor_hr, w, target_lr_y_means[i])
        predicted_patches.append(pred_patch)
        
    # 3. Ghép các patch lại thành ảnh (Aggregation)
    lr_img = cv2.imread(lr_img_path)
    h_lr, w_lr = lr_img.shape[:2]
    h_hr, w_hr = h_lr * scale, w_lr * scale
    
    predicted_y = np.zeros((h_hr, w_hr), dtype=np.float32)
    weight_mask = np.zeros((h_hr, w_hr), dtype=np.float32)
    
    step = int(patch_size * (1 - overlap))
    patch_idx = 0
    
    for i in range(0, h_lr - patch_size + 1, step):
        for j in range(0, w_lr - patch_size + 1, step):
            if patch_idx >= len(predicted_patches): break
            
            hr_i, hr_j = i * scale, j * scale
            p_h, p_w = predicted_patches[patch_idx].shape
            
            predicted_y[hr_i:hr_i+p_h, hr_j:hr_j+p_w] += predicted_patches[patch_idx]
            weight_mask[hr_i:hr_i+p_h, hr_j:hr_j+p_w] += 1.0
            patch_idx += 1
            
    # Normalize (chia trung bình vùng chồng lấn)
    weight_mask[weight_mask == 0] = 1.0
    predicted_y /= weight_mask
    
    return predicted_y * 255.0  # Trả về range 0-255

def main():
    # --- CẤU HÌNH ---
    parser = argparse.ArgumentParser(description="Evaluate Neighbor Embedding SR")
    parser.add_argument('--scale', type=int, default=4, help='Upscale factor (2, 3, 4)')
    parser.add_argument('--patch_size', type=int, default=5, help='Patch size')
    parser.add_argument('--k', type=int, default=5, help='Number of neighbors for KNN')
    
    args = parser.parse_args()
    
    SCALE = args.scale
    PATCH_SIZE = args.patch_size
    OVERLAP = 0.5
    K_NEIGHBORS = args.k
    EPSILON = 1e-6

    # Đường dẫn dữ liệu
    DATA_ROOT = '../data' # Sửa lại nếu cần
    TRAIN_FEATS_PATH = os.path.join(DATA_ROOT, f'YIQ_patches/LR_features_patches_x{SCALE}.npy')
    TRAIN_LABEL_PATH = os.path.join(DATA_ROOT, f'YIQ_patches/HR_centered_patches_x{SCALE}.npy')
    TEST_LR_FOLDER = os.path.join(DATA_ROOT, f'target_lr/LR_x{SCALE}')
    TEST_HR_FOLDER = os.path.join(DATA_ROOT, f'target_hr/HR_x{SCALE}')

    print(f"=== BẮT ĐẦU ĐÁNH GIÁ NE (Scale x{SCALE}) ===")
    
    # 1. Load dữ liệu Training
    if not os.path.exists(TRAIN_FEATS_PATH):
        print(f"Error: Không tìm thấy file {TRAIN_FEATS_PATH}")
        return

    print("-> Loading training patches...")
    LR_features_patches = np.load(TRAIN_FEATS_PATH)
    HR_centered_patches = np.load(TRAIN_LABEL_PATH)
    
    # 2. Build KNN
    print(f"-> Building KNN Index ({LR_features_patches.shape[0]} patches)...")
    t0 = time.time()
    knn_model, feat_vectors = ml_patch.build_knn_idx(LR_features_patches, k=K_NEIGHBORS)
    print(f"-> KNN Built in {time.time() - t0:.2f}s")

    # 3. Chuẩn bị Test Set
    test_images = sorted(glob(os.path.join(TEST_LR_FOLDER, '*.png')))
    gt_images = sorted(glob(os.path.join(TEST_HR_FOLDER, '*.png')))
    
    if not test_images:
        print("Error: Không tìm thấy ảnh test trong folder Set5.")
        return

    # 4. Vòng lặp đánh giá
    print("\n" + "="*85)
    print(f"{'Image Name':<20} | {'Bicubic (PSNR/SSIM)':<25} | {'NE (PSNR/SSIM)':<25} | {'Gain':<8}")
    print("-" * 85)

    results = [] # [psnr_bic, ssim_bic, psnr_ne, ssim_ne]

    for lr_path, hr_path in zip(test_images, gt_images):
        img_name = os.path.basename(lr_path)
        
        # -- A. Ground Truth (HR) --
        hr_img = cv2.imread(hr_path)
        hr_img = cv2.cvtColor(hr_img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        hr_yiq = ml_patch.rgb2yiq(hr_img)
        hr_y = hr_yiq[:, :, 0] * 255.0
        
        # -- B. Bicubic Baseline --
        lr_img = cv2.imread(lr_path)
        lr_img = cv2.cvtColor(lr_img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        h, w = lr_img.shape[:2]
        bic_hr = cv2.resize(lr_img, (w * SCALE, h * SCALE), interpolation=cv2.INTER_CUBIC)
        bic_yiq = ml_patch.rgb2yiq(bic_hr)
        bic_y = bic_yiq[:, :, 0] * 255.0
        
        # -- C. NE Reconstruction --
        ne_y = run_ne_inference(lr_path, knn_model, feat_vectors, HR_centered_patches, 
                                SCALE, PATCH_SIZE, OVERLAP, K_NEIGHBORS, EPSILON)
        
        # Clip giá trị hợp lệ (0-255)
        ne_y = np.clip(ne_y, 0, 255)
        bic_y = np.clip(bic_y, 0, 255)
        
        # -- D. Tính Metrics (PSNR/SSIM trên kênh Y, có cắt viền) --
        border_cut = SCALE*PATCH_SIZE // 2
        p_bic = utils.calculate_psnr(bic_y, hr_y, crop_border=border_cut)
        s_bic = utils.calculate_ssim(bic_y, hr_y, crop_border=border_cut)
        
        p_ne = utils.calculate_psnr(ne_y, hr_y, crop_border=border_cut)
        s_ne = utils.calculate_ssim(ne_y, hr_y, crop_border=border_cut)
        
        results.append([p_bic, s_bic, p_ne, s_ne])
        
        print(f"{img_name:<20} | {p_bic:.2f} dB / {s_bic:.4f}        | {p_ne:.2f} dB / {s_ne:.4f}        | {p_ne - p_bic:+.2f}")

    # 5. Tính trung bình
    avg = np.mean(results, axis=0)
    print("-" * 85)
    print(f"{'AVERAGE':<20} | {avg[0]:.2f} dB / {avg[1]:.4f}        | {avg[2]:.2f} dB / {avg[3]:.4f}        | {avg[2] - avg[0]:+.2f}")
    print("=" * 85)

if __name__ == "__main__":
    main()