Spaces:
Running
Running
| 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" | |
| RESNET_MODEL_FILENAME = "Lib/Models/Image/resnet_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) | |
| resnet_model_path = hf_hub_download(repo_id=REPO_ID, filename=RESNET_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(clip_model_path) | |
| cnn_model = tf.keras.models.load_model(cnn_model_path) | |
| resnet_model = tf.keras.models.load_model(resnet_model_path) | |
| 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"skipped image (too small) ({w}x{h})" | |
| 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 center: {e}" | |
| def denoise(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 = 10 # filter strength for luminance component (recommended 3-15) | |
| H_COLOR = 10 # 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 | |
| ) | |
| #cv2.imwrite("denoised.png", den) # for debugging | |
| den = cv2.cvtColor(den, cv2.COLOR_BGR2RGB) | |
| den = Image.fromarray(den) | |
| return den | |
| def compute_profile(raw_image: Image, den_image: Image, normalize=False ,verbose= True) -> np.ndarray | str: | |
| # read images | |
| if raw_image is None: | |
| return print(f"WARNING: couldn't read raw image") | |
| if den_image is None: | |
| return print(f"WARNING: couldn't read denoised image") | |
| raw = np.array(raw_image) # RGB uint8 numpy array | |
| raw = cv2.cvtColor(raw, cv2.COLOR_RGB2BGR) | |
| den = np.array(den_image) # RGB uint8 numpy array | |
| den = cv2.cvtColor(den, cv2.COLOR_RGB2BGR) | |
| # if shapes differ, resize den to raw's size (keeps alignment); warn | |
| if den.shape != raw.shape: | |
| if verbose: | |
| print(f"NOTE: shape mismatch, resizing denoised from {den.shape[:2]} to {raw.shape[:2]}") | |
| den = cv2.resize(den, (raw.shape[1], raw.shape[0]), interpolation=cv2.INTER_LINEAR) | |
| # absolute difference per-channel | |
| diff = cv2.absdiff(raw, den) # BGR, uint8 | |
| gray = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY) # single-channel uint8 | |
| # optionally normalize to full 0-255 (per-image) | |
| if normalize: | |
| # cv2.normalize will map min->0 and max->255 | |
| # but if the image is flat (min==max) normalize will set to 0; handle that | |
| 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 | |
| else: | |
| # keep raw diff values but ensure dtype uint8 (already uint8) and values are 0..255 | |
| out = gray | |
| #cv2.imwrite("profile.png", out) # for debugging | |
| return out | |
| def profile_image_for_cnn_predict(pil_img: Image, crop_size=512): | |
| """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, crop_size=crop_size) | |
| if isinstance(cropped_img, str): | |
| return cropped_img # return error message if cropping failed | |
| # Step 2: Denoise the cropped image | |
| denoised_img = denoise(cropped_img) | |
| if isinstance(denoised_img, str): | |
| return denoised_img # return error message if denoising failed | |
| # Step 3: Compute the profile image | |
| profile_img = compute_profile(cropped_img, denoised_img, normalize=False) | |
| if isinstance(profile_img, str): | |
| return profile_img # return error message if profile computation failed | |
| return profile_img | |
| def prepare_cv2_image_for_resnet(cv2_gray_img, target_size=(512,512)): | |
| img_rgb = cv2.cvtColor(cv2_gray_img, cv2.COLOR_GRAY2RGB) | |
| img_rgb = cv2.resize(img_rgb, (target_size[1], target_size[0]), interpolation=cv2.INTER_AREA) | |
| img_rgb = img_rgb.astype('float32') | |
| # 5) add batch dim | |
| x = np.expand_dims(img_rgb, axis=0) # shape (1, H, W, 3) | |
| x = preprocess_input(x) | |
| return x | |
| def predict_image_prob_clip(image: Image.Image, threshold=0.5, | |
| clip_model=None, clip_preprocess=None, | |
| keras_mlp=None): | |
| """ | |
| 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'} | |
| """ | |
| # --- 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: | |
| print("No keras model provided...") | |
| return None | |
| # --- load/normalize image --- | |
| # assume PIL image | |
| 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 clip_predict(pil_img: Image, crop_size=512): | |
| # pass model objects explicitly (faster if you call this repeatedly) | |
| pil_img = center_crop(pil_img, crop_size=crop_size) | |
| if isinstance(pil_img, str): | |
| return pil_img # return error message | |
| return predict_image_prob_clip(pil_img, | |
| clip_model=clip_mod, | |
| clip_preprocess=clip_pre, | |
| keras_mlp=mlp_model) | |
| def CNNPredict(predict_img: np.ndarray): | |
| #1 Real 0 AI | |
| #normalize image | |
| # expand dims to add channel axis | |
| 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] | |
| def ResnetPredict(predict_img): | |
| #1 Real 0 AI | |
| predict_img = prepare_cv2_image_for_resnet(predict_img) | |
| prediction = resnet_model.predict(predict_img) | |
| return prediction[0][0] | |