ClassiferTest.py

"""
Script doc lap de test model phan loai u nao da train (.keras).
Khong can toan bo pipeline training - chi can model + 1 file de test.

Ho tro 2 loai input:
  1. File .h5 (dung dinh dang giong data training, co dataset 'image')
  2. File anh thong thuong (.png, .jpg, .jpeg) - se duoc chuyen ve grayscale

Cach dung:
    python test_inference.py --model /path/to/model.keras --input /path/to/slice.h5
    python test_inference.py --model /path/to/model.keras --input /path/to/image.png
"""

import argparse
import numpy as np
import cv2
import h5py
import tensorflow as tf
from tensorflow.keras.layers import Layer, Dense, Dropout
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input as efficientnetv2_preprocess

IMG_SIZE = (299, 299)
NUM_SLICES = 3
CLASS_NAMES = ["non-tumor", "tumor"]


# ============================================================
# CAC CUSTOM LAYER / LOSS - BAT BUOC PHAI CO DE load_model() HOAT DONG
# (copy y nguyen tu script training goc)
# ============================================================
class AttentionVisualizer(Layer):
    def __init__(self, **kwargs):
        super(AttentionVisualizer, self).__init__(**kwargs)

    def call(self, attention_weights):
        spatial_attention = tf.reduce_mean(attention_weights, axis=-1, keepdims=True)
        spatial_attention = (spatial_attention - tf.reduce_min(spatial_attention, axis=(1, 2), keepdims=True)) / \
                             (tf.reduce_max(spatial_attention, axis=(1, 2), keepdims=True) -
                              tf.reduce_min(spatial_attention, axis=(1, 2), keepdims=True) + 1e-8)
        return spatial_attention


class MCDropout(Layer):
    def __init__(self, rate, **kwargs):
        super(MCDropout, self).__init__(**kwargs)
        self.rate = rate
        self.dropout = Dropout(rate)

    def call(self, inputs, training=None):
        return self.dropout(inputs, training=True)


class EvidentialLoss(tf.keras.losses.Loss):
    def __init__(self, class_weights=None, name="evidential_loss",
                 reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE):
        super(EvidentialLoss, self).__init__(name=name, reduction=reduction)
        self.class_weights = class_weights if class_weights is not None else {}

    def call(self, y_true, y_pred):
        ce_loss = tf.keras.losses.categorical_crossentropy(y_true, y_pred)
        if self.class_weights:
            class_indices = tf.argmax(y_true, axis=1)
            weights = tf.gather(tf.constant(list(self.class_weights.values()), dtype=tf.float32), class_indices)
            ce_loss = ce_loss * weights
        reg_loss = tf.reduce_mean(tf.square(y_pred - y_true))
        return tf.reduce_mean(ce_loss + 0.01 * reg_loss)

    def get_config(self):
        config = super(EvidentialLoss, self).get_config()
        config.update({'class_weights': self.class_weights})
        return config


class EvidentialLayer(Layer):
    def __init__(self, num_classes, **kwargs):
        super(EvidentialLayer, self).__init__(**kwargs)
        self.num_classes = num_classes

    def build(self, input_shape):
        self.dense = Dense(self.num_classes, dtype=tf.float32)
        super(EvidentialLayer, self).build(input_shape)

    def call(self, inputs):
        evidence = tf.nn.softplus(self.dense(inputs))
        evidence = tf.clip_by_value(evidence, 0, 1e6)
        alpha = evidence + 1
        S = tf.reduce_sum(alpha, axis=1, keepdims=True) + 1e-10
        prob = alpha / S
        epistemic_uncertainty = self.num_classes / S
        aleatoric_uncertainty = tf.reduce_sum(prob * (1 - prob) / (S + 1), axis=1, keepdims=True)
        return prob, epistemic_uncertainty, aleatoric_uncertainty

    def get_config(self):
        config = super(EvidentialLayer, self).get_config()
        config.update({'num_classes': self.num_classes})
        return config


class ExplainableSelfAttention(Layer):
    def __init__(self, units, return_attention=False, **kwargs):
        super(ExplainableSelfAttention, self).__init__(**kwargs)
        self.units = units
        self.return_attention = return_attention

    def build(self, input_shape):
        self.W_q = Dense(self.units, use_bias=False)
        self.W_k = Dense(self.units, use_bias=False)
        self.W_v = Dense(self.units, use_bias=False)
        self.dense = Dense(input_shape[-1])
        self.attention_visualizer = AttentionVisualizer()
        super(ExplainableSelfAttention, self).build(input_shape)

    def call(self, inputs, return_attention=None):
        return_attention = return_attention if return_attention is not None else self.return_attention
        batch_size = tf.shape(inputs)[0]
        height, width, channels = inputs.shape[1:]
        x = tf.reshape(inputs, [batch_size, height * width, channels])
        q = self.W_q(x)
        k = self.W_k(x)
        v = self.W_v(x)
        scale = tf.cast(tf.sqrt(tf.cast(self.units, inputs.dtype)), inputs.dtype)
        scores = tf.matmul(q, k, transpose_b=True) / scale
        attention_weights = tf.nn.softmax(scores, axis=-1)
        attended = tf.matmul(attention_weights, v)
        output = self.dense(attended)
        output = tf.reshape(output, [batch_size, height, width, channels])
        output = inputs + output
        if return_attention:
            attention_spatial = tf.reshape(attention_weights, [batch_size, height, width, height * width])
            attention_map = self.attention_visualizer(attention_spatial)
            return output, attention_map
        return output

    def get_config(self):
        config = super(ExplainableSelfAttention, self).get_config()
        config.update({'units': self.units, 'return_attention': self.return_attention})
        return config


class MaxProbLayer(Layer):
    def __init__(self, **kwargs):
        super(MaxProbLayer, self).__init__(**kwargs)

    def call(self, inputs):
        return tf.reduce_max(inputs, axis=1, keepdims=True)


CUSTOM_OBJECTS = {
    'AttentionVisualizer': AttentionVisualizer,
    'MCDropout': MCDropout,
    'EvidentialLoss': EvidentialLoss,
    'EvidentialLayer': EvidentialLayer,
    'ExplainableSelfAttention': ExplainableSelfAttention,
    'MaxProbLayer': MaxProbLayer,
}


# ============================================================
# TIEN XU LY ANH - GIONG HET LOGIC TRONG GENERATOR TRAINING
# ============================================================
def load_slice_from_h5(path):
    """Doc 1 slice tu file .h5 (dataset ten 'image')."""
    with h5py.File(path, 'r') as f:
        if 'image' not in f:
            raise ValueError(f"File h5 khong co dataset 'image': {path}")
        img = f['image'][:].astype(np.float32)
        if img.ndim == 3:
            img = img[..., 0]
        elif img.ndim != 2:
            raise ValueError(f"Shape anh khong hop le: {img.shape}")
        return img


def load_slice_from_image(path):
    """Doc 1 slice tu file anh thuong (.png/.jpg), chuyen ve grayscale."""
    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise ValueError(f"Khong doc duoc anh: {path}")
    return img.astype(np.float32)


def preprocess_slice(img_slice):
    """Resize, chuan hoa mean/std, clip, tao 3-channel stack, roi ap dung
    efficientnetv2_preprocess - dung y nhu trong _load_slice/_generate_batch
    cua ImprovedBrainTumorGenerator ban goc."""
    if img_slice.shape != IMG_SIZE:
        img_slice = cv2.resize(img_slice, IMG_SIZE, interpolation=cv2.INTER_LINEAR)

    mean = np.mean(img_slice)
    std = np.std(img_slice) + 1e-8
    img_slice = (img_slice - mean) / std
    img_slice = np.clip(img_slice, -3, 3)

    # gia lap 3 slice lan can bang cach lap lai slice trung tam
    stack = np.stack([img_slice, img_slice, img_slice], axis=-1)

    X = np.expand_dims(stack, axis=0).astype(np.float32)  # (1, 299, 299, 3)
    X = efficientnetv2_preprocess(X)
    return X


def load_and_preprocess(path):
    if path.lower().endswith(('.h5', '.hdf5')):
        img_slice = load_slice_from_h5(path)
    else:
        img_slice = load_slice_from_image(path)
    return preprocess_slice(img_slice)


# ============================================================
# CHAY INFERENCE
# ============================================================
def run_inference(model_path, input_path):
    print(f"Dang load model: {model_path}")
    model = tf.keras.models.load_model(model_path, custom_objects=CUSTOM_OBJECTS)

    print(f"Dang doc va tien xu ly file: {input_path}")
    X = load_and_preprocess(input_path)

    print("Dang du doan...")
    outputs = model.predict(X, verbose=0)

    # model tra ve dict cac output (classification, evidential_prob, confidence_score, ...)
    probs = outputs['classification'][0]
    pred_idx = int(np.argmax(probs))
    pred_label = CLASS_NAMES[pred_idx]

    print("\n" + "=" * 50)
    print("KET QUA DU DOAN")
    print("=" * 50)
    print(f"Nhan du doan: {pred_label}")
    for i, name in enumerate(CLASS_NAMES):
        print(f"  P({name}) = {probs[i]:.4f}")

    if 'confidence_score' in outputs:
        print(f"Confidence score: {float(outputs['confidence_score'][0][0]):.4f}")
    if 'epistemic_uncertainty' in outputs:
        print(f"Epistemic uncertainty: {float(outputs['epistemic_uncertainty'][0][0]):.4f}")
    if 'aleatoric_uncertainty' in outputs:
        print(f"Aleatoric uncertainty: {float(outputs['aleatoric_uncertainty'][0][0]):.4f}")
    print("=" * 50)

    return outputs


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test model phan loai u nao voi 1 file don le")
    parser.add_argument("--model", required=True, help="Duong dan file model .keras")
    parser.add_argument("--input", required=True, help="Duong dan file .h5 hoac anh (.png/.jpg) can test")
    args = parser.parse_args()

    run_inference(args.model, args.input)