Spaces:
Running on Zero
Running on Zero
File size: 5,098 Bytes
b5aa172 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | import numpy as np
import tensorflow as tf
from PIL import Image
import json
import os
IMG_SIZE = 300
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
model_dir = os.path.join(BASE_DIR, "final_model_defactify")
def _sanitize_config(obj):
if isinstance(obj, dict):
obj.pop("quantization_config", None)
for v in obj.values():
_sanitize_config(v)
elif isinstance(obj, list):
for item in obj:
_sanitize_config(item)
return obj
def load_model():
config_path = os.path.join(model_dir, "config.json")
weights_path = os.path.join(model_dir, "model.weights.h5")
if not os.path.exists(config_path):
raise FileNotFoundError(f"Model config not found at: {config_path}")
with open(config_path, "r") as f:
config = json.load(f)
_sanitize_config(config)
grad_cam_model = tf.keras.models.model_from_json(
json.dumps(config)
)
if not os.path.exists(weights_path):
raise FileNotFoundError(
f"Weights file not found at '{weights_path}'. "
"Please ensure 'model.weights.h5' is placed in 'models/final_model_defactify/'."
)
# Load trained weights
grad_cam_model.load_weights(weights_path)
return grad_cam_model
def preprocess_image(image: Image.Image) -> np.ndarray:
image = image.convert("RGB")
image = image.resize((IMG_SIZE, IMG_SIZE))
arr = np.asarray(image, dtype=np.float32) # (IMG_SIZE, IMG_SIZE, 3), [0,255]
arr = np.expand_dims(arr, axis=0) # add batch dim
arr = arr.astype(np.float32) # match interpreter's expected dtype
return arr
def find_submodel(model):
"""Find the nested backbone sub-model (e.g. EfficientNetV2B3) inside
the outer classification model."""
for layer in model.layers:
if isinstance(layer, tf.keras.Model):
return layer
raise ValueError("No nested sub-model (backbone) found inside grad_cam_model.")
def find_last_conv_layer(model):
"""Find the last conv-like (4D output) layer, recursing into nested models."""
for layer in reversed(model.layers):
if isinstance(layer, tf.keras.Model):
try:
return find_last_conv_layer(layer)
except ValueError:
continue
if len(layer.output.shape) == 4: # (batch, H, W, channels)
return layer.name
raise ValueError("No 4D conv-like layer found.")
def make_gradcam_heatmap(image_array,model):
"""image_array: output of preprocess_image() -- (1, IMG_SIZE, IMG_SIZE, 3), [0,255]."""
base_model = find_submodel(model)
last_conv_layer_name = find_last_conv_layer(base_model)
conv_layer_model = tf.keras.Model(
base_model.input, base_model.get_layer(last_conv_layer_name).output
)
# Rebuild the head (GAP -> Dropout -> Dense) as a separate model so we
# can run it on top of the conv output inside the same GradientTape.
head_input = tf.keras.Input(shape=conv_layer_model.output.shape[1:])
x = head_input
for layer in model.layers:
if isinstance(layer, tf.keras.Model) or isinstance(layer, tf.keras.layers.InputLayer):
continue
x = layer(x)
head_model = tf.keras.Model(head_input, x)
image_tensor = tf.convert_to_tensor(image_array, dtype=tf.float32)
with tf.GradientTape() as tape:
conv_output = conv_layer_model(image_tensor)
tape.watch(conv_output)
preds = head_model(conv_output)
class_channel = preds[:, 0] # sigmoid output, single unit
grads = tape.gradient(class_channel, conv_output)
pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
conv_output = conv_output[0]
heatmap = conv_output @ pooled_grads[..., tf.newaxis]
heatmap = tf.squeeze(heatmap)
heatmap = tf.maximum(heatmap, 0) / (tf.reduce_max(heatmap) + 1e-8)
return heatmap.numpy(), float(preds[0, 0])
def predict(image: Image.Image,model , alpha=0.4):
"""Takes a raw PIL image, runs preprocessing + Grad-CAM, shows the overlay."""
image_array = preprocess_image(image) # reuse the same preprocessing as deployment
image_array = image_array.astype(np.float32) # ensure float32 for the gradient model
heatmap, prob = make_gradcam_heatmap(image_array,model)
heatmap_resized = tf.image.resize(heatmap[..., tf.newaxis], [IMG_SIZE, IMG_SIZE]).numpy().squeeze()
r = np.clip(1.5 - np.abs(4 * heatmap_resized - 3), 0, 1)
g = np.clip(1.5 - np.abs(4 * heatmap_resized - 2), 0, 1)
b = np.clip(1.5 - np.abs(4 * heatmap_resized - 1), 0, 1)
heatmap_colored = np.stack([r, g, b], axis=-1)
orig = image_array[0] / 255.0
overlay = orig * (1 - alpha) + heatmap_colored * alpha
pil_overlay = np.clip(overlay * 255, 0, 255).astype(np.uint8)
pil_overlay = Image.fromarray(pil_overlay)
label = "AI Generated Image" if prob > 0.5 else "Real Image"
confidence = prob if prob > 0.5 else 1 - prob
print( label , confidence )
print(type(pil_overlay))
return label,confidence,pil_overlay |