imagePreprocess.py

import os
from pathlib import Path
from PIL import Image, ImageOps
import cv2
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.models import load_model
import torch
import clip
from huggingface_hub import hf_hub_download

BASE_DIR = "MAS-AI-0000/Authentica"
MODELS_DIR = os.path.join(BASE_DIR, "Lib/Models/Image")

# ==== CONFIG ====
REPO_ID = "MAS-AI-0000/Authentica"
CLIP_MODEL_FILENAME = "Lib/Models/Image/clip_model.keras"
CNN_MODEL_FILENAME = "Lib/Models/Image/cnn_model.keras"

# ==== Load assets ====
clip_model_path = hf_hub_download(repo_id=REPO_ID, filename=CLIP_MODEL_FILENAME)
cnn_model_path = hf_hub_download(repo_id=REPO_ID, filename=CNN_MODEL_FILENAME)

# Load models and preprocessing once at module level
clip_mod, clip_pre = clip.load("ViT-B/32", jit=False)
clip_mod.eval()
for p in clip_mod.parameters():
    p.requires_grad = False

mlp_model= tf.keras.models.load_model(os.path.join(MODELS_DIR, "clip_model.keras"))
cnn_model = tf.keras.models.load_model(os.path.join(MODELS_DIR, "cnn_model.keras"))


def center_crop(image: Image.Image, crop_size=512) -> Image.Image | str:
    try:
            image = ImageOps.exif_transpose(image)
            w, h = image.size
            if w < crop_size or h < crop_size:
                # skip small images
                return f"Image is too small: ({w}x{h}), Minimum size is {crop_size}x{crop_size}"
            left = (w - crop_size) // 2
            top = (h - crop_size) // 2
            right = left + crop_size
            bottom = top + crop_size
            cropped = image.crop((left, top, right, bottom))
            return cropped
    except Exception as e:
        return f"Error when cropping image: {e}"


def compute_profile(src_image: Image) -> np.ndarray | str:
    """Read image, denoise (GPU if available) and return denoised image."""
    img = np.array(src_image)   # BGR uint8 numpy array
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    if src_image is None:
        print(f"WARNING: No source image, skipping.")
        return False
    # Denoising parameters
    H = 5           # filter strength for luminance component (recommended 3-15)
    H_COLOR = 5     # same for color components
    TEMPLATE_WINDOW_SIZE = 7
    SEARCH_WINDOW_SIZE = 21
    # Use CUDA if available, otherwise CPU fallback
    use_cuda = False
    try:
        use_cuda = hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0
    except Exception:
        use_cuda = False
    if use_cuda:
        # Create a GpuMat and upload the numpy image to GPU
        gpu_img = cv2.cuda_GpuMat()
        gpu_img.upload(img)   # <-- this converts numpy -> GpuMat on device
        den_gpu = cv2.cuda.fastNlMeansDenoisingColored(
           gpu_img,H,H_COLOR,None,SEARCH_WINDOW_SIZE,TEMPLATE_WINDOW_SIZE
        )

        # Download result back to CPU
        den = den_gpu.download()
    else:
        # Fallback to CPU implementation
        print("NOTICE: CUDA not available — using CPU denoiser.")
        den = cv2.fastNlMeansDenoisingColored(
            img, None,
            H, H_COLOR,
            TEMPLATE_WINDOW_SIZE,
            SEARCH_WINDOW_SIZE
        )

    # absolute difference per-channel
    diff = cv2.absdiff(img, den)               # BGR, uint8
    gray = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)  # single-channel uint8

    minv = int(gray.min())
    maxv = int(gray.max())
    if maxv > minv:
            norm = cv2.normalize(gray, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
            out = norm
    else:
        # nothing to normalize (flat), keep as-is (all zeros)
        out = gray
    return out

def preprocess_cnn(pil_img: Image):
    """Preprocess the input image and return a numpy array ready for model prediction."""
    # Step 1: Center crop the image
    cropped_img = center_crop(pil_img)
    if isinstance(cropped_img, str):
        return cropped_img  # return error message if cropping failed
    # Step 3: Compute the profile image
    profile_img = compute_profile(cropped_img)
    if isinstance(profile_img, str):
        return profile_img  # return error message if profile computation failed
    return profile_img


def CLIPPredict(image: Image.Image, 
                       clip_model=clip_mod, clip_preprocess=clip_pre,
                       keras_mlp=mlp_model) -> float | str:
    """
    Predicts probability that image is AI-generated (AI=1) using CLIP + Keras MLP.

    Args:
      path_or_image: str (file path) or PIL.Image.Image or numpy array (H,W,3)
      threshold: float threshold for binary label
      clip_model, clip_preprocess: optionally pass existing CLIP objects
      keras_mlp: optionally pass existing loaded Keras model
    Returns:
      dict: {'prob': float_prob_AI, 'label': 'AI' or 'Real'}
    """
    #0 Real 1 AI
    # --- try to reuse provided CLIP objects, otherwise load ---
    if clip_model is None or clip_preprocess is None:
        print("Loading Default CLIP model...")
        # pick a model name: prefer provided arg, else try global, else ViT-B/32
        cmn = "ViT-B/32"
        clip_model, clip_preprocess = clip.load(cmn, device="cpu", jit=False)
        clip_model.eval()
        for p in clip_model.parameters():
            p.requires_grad = False

    # --- try to reuse provided keras model, otherwise load from disk ---
    if keras_mlp is None:
       return  "No keras model provided..."
    # --- load/normalize image ---
    # assume PIL image
    image = center_crop(image, crop_size=512)
    if isinstance(image, str):
        return image  # return error message if cropping failed
    img = image.convert('RGB')
   
    # --- preprocess for CLIP and get embedding ---
    input_tensor = clip_preprocess(img).unsqueeze(0).to("cpu")  # shape (1,C,H,W)
    with torch.no_grad():
        emb = clip_model.encode_image(input_tensor)   # (1, D)
        emb = emb / emb.norm(dim=-1, keepdim=True)    # L2 normalize

    emb_np = emb.cpu().numpy().astype('float32')     # shape (1, D)

    # --- predict with Keras MLP ---
    probs = keras_mlp.predict(emb_np, verbose=0).reshape(-1,)
    prob = float(probs[0])
    return prob


def CNNPredict(img: Image.Image) -> float | str:
    predict_img = preprocess_cnn(img)
    if isinstance(predict_img, str):
        return predict_img  # return error message if preprocessing failed
    predict_img = predict_img.astype('float32') / 255.0  # shape (H, W)
    predict_img = np.expand_dims(predict_img, axis=-1)  # shape (H, W, 1)
    # expand dims to add batch axis
    predict_img = np.expand_dims(predict_img, axis=0)   # shape (1, H, W, 1)
    prediction = cnn_model.predict(predict_img)
    return prediction[0][0]