import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import numpy as np

# --- Configuration ---
MODEL_PATH = 'weapon_classifier_final_tuned.keras'
IMAGE_PATH = './test_image.jpg' 
IMG_SIZE = (224, 224)
# --- End Configuration ---

def load_and_preprocess_image(img_path, target_size):
    """Loads, resizes, and normalizes the image for prediction."""
    img = image.load_img(img_path, target_size=target_size)
    img_array = image.img_to_array(img)
    # Add batch dimension: (H, W, C) -> (1, H, W, C)
    img_array = np.expand_dims(img_array, axis=0)
    # Normalize pixel values (0-255 -> 0-1)
    processed_image = img_array / 255.0
    return processed_image

def classify_image(model_path, image_path, img_size):
    """Loads the model, makes a prediction, and interprets the result."""
    
    # Load the model
    model = load_model(model_path)
    print("Model loaded successfully.")

    # Preprocess the image
    input_image = load_and_preprocess_image(image_path, img_size)

    # Make the prediction
    prediction = model.predict(input_image)

    # Interpret the result for binary classification
    probability = prediction[0][0]
    class_names = {0: 'Not a Weapon', 1: 'Weapon'}

    if probability >= 0.5:
        predicted_class = class_names[1]
        confidence = probability * 100
    else:
        predicted_class = class_names[0]
        confidence = (1 - probability) * 100 

    print("\n--- CLASSIFICATION RESULT ---")
    print(f"Image: {os.path.basename(image_path)}")
    print(f"Predicted Class: **{predicted_class}**")
    print(f"Confidence: **{confidence:.2f}%**")
    print("---------------------------\n")

# --- EXECUTE CLASSIFICATION ---
import os
classify_image(MODEL_PATH, IMAGE_PATH, IMG_SIZE)