import os import re import io import random import numpy as np import tensorflow as tf import cv2 from PIL import Image, ImageChops, ImageDraw from sklearn.model_selection import train_test_split from tensorflow.keras import layers, models, applications # ── Global configuration ────────────────────────────────────────────────────── SEED = 42 IMG_SIZE = (224, 224) ELA_QUALITY = 90 ELA_SCALE = 15 BATCH_SIZE = 32 EPOCHS = 5 TARGET_DIR = "./casia_v2" def set_reproducibility(seed=SEED): tf.random.set_seed(seed) np.random.seed(seed) random.seed(seed) os.environ['PYTHONHASHSEED'] = str(seed) set_reproducibility() def generate_robust_dataset(num_samples=120): if os.path.exists(TARGET_DIR): import shutil shutil.rmtree(TARGET_DIR) os.makedirs(TARGET_DIR) print(f"Generating {num_samples} synthetic samples...") for i in range(num_samples): img_data = np.random.randint(100, 200, (256, 256, 3), dtype=np.uint8) img = Image.fromarray(img_data) is_forged = i >= (num_samples // 2) if not is_forged: filename = f"Au_arc_000{i:02d}.jpg" else: draw = ImageDraw.Draw(img) draw.rectangle([50, 50, 150, 150], fill=(255, 0, 0)) filename = f"Tp_s_N_arc_000{i:02d}_00099_001.jpg" img.save(os.path.join(TARGET_DIR, filename)) def compute_ela(image_path_or_pil, quality=ELA_QUALITY, scale=ELA_SCALE): if isinstance(image_path_or_pil, str): original = Image.open(image_path_or_pil).convert('RGB') else: original = image_path_or_pil.convert('RGB') buf = io.BytesIO() original.save(buf, 'JPEG', quality=quality) buf.seek(0) compressed = Image.open(buf) ela_image = ImageChops.difference(original, compressed) ela_image = ImageChops.multiply( ela_image, Image.new('RGB', ela_image.size, (scale, scale, scale)) ) return ela_image class CASIAParser: @staticmethod def get_ids(filename): name = os.path.basename(filename) if name.startswith('Au_'): match = re.search(r'Au_[a-z]{3}_(\d+)', name) return [match.group(1)] if match else [] elif name.startswith('Tp_'): parts = name.split('_') return [parts[4], parts[5]] if len(parts) >= 6 else [] return [] def split_dataset(data_dir, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1): all_images = [ os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.tif')) ] unique_ids = sorted({i for p in all_images for i in CASIAParser.get_ids(p)}) if not unique_ids: unique_ids = [str(i) for i in range(len(all_images))] tr_ids, temp = train_test_split(unique_ids, train_size=train_ratio, random_state=SEED) v_ids, _ = train_test_split(temp, train_size=val_ratio / (val_ratio + test_ratio), random_state=SEED) tr_ids, v_ids = set(tr_ids), set(v_ids) splits = {'train': [], 'val': [], 'test': []} for p in all_images: ids = CASIAParser.get_ids(p) if not ids: splits['train'].append(p) if random.random() < 0.8 else splits['test'].append(p) continue if any(i in tr_ids for i in ids): splits['train'].append(p) elif any(i in v_ids for i in ids): splits['val'].append(p) else: splits['test'].append(p) return splits def preload_images(paths, img_size=IMG_SIZE): rgb_list, ela_list, label_list = [], [], [] for p in paths: pil_img = Image.open(p).convert('RGB') rgb_list.append(np.array(pil_img.resize(img_size), dtype=np.float32)) ela_list.append(np.array(compute_ela(pil_img).resize(img_size), dtype=np.float32)) label_list.append(1 if os.path.basename(p).startswith('Tp_') else 0) return np.array(rgb_list), np.array(ela_list), np.array(label_list) def make_dataset(rgb_arr, ela_arr, labels, batch_size=BATCH_SIZE, shuffle=False, repeat=True): ds = tf.data.Dataset.from_tensor_slices(((rgb_arr, ela_arr), labels)) if shuffle: ds = ds.shuffle(buffer_size=len(labels), seed=SEED, reshuffle_each_iteration=True) ds = ds.batch(batch_size, drop_remainder=False) if repeat: ds = ds.repeat() return ds.prefetch(tf.data.AUTOTUNE) def get_rgb_branch(): base = applications.ResNet50( include_top=False, weights='imagenet', input_shape=(*IMG_SIZE, 3) ) base.trainable = False inputs = layers.Input(shape=(*IMG_SIZE, 3)) x = applications.resnet50.preprocess_input(inputs) x = base(x, training=False) return inputs, layers.GlobalAveragePooling2D()(x) def get_ela_branch(): inputs = layers.Input(shape=(*IMG_SIZE, 3)) x = layers.Rescaling(1. / 255)(inputs) for filters in [32, 64, 128]: x = layers.Conv2D(filters, (3, 3), activation='relu', padding='same')(x) x = layers.BatchNormalization()(x) x = layers.MaxPooling2D((2, 2))(x) return inputs, layers.GlobalAveragePooling2D()(x) def build_model(): rgb_in, rgb_f = get_rgb_branch() ela_in, ela_f = get_ela_branch() fused = layers.Concatenate()([rgb_f, ela_f]) out = layers.Dense(1, activation='sigmoid')( layers.Dropout(0.5)(layers.Dense(256, activation='relu')(fused)) ) return models.Model(inputs=[rgb_in, ela_in], outputs=out) if __name__ == "__main__": generate_robust_dataset(120) splits = split_dataset(TARGET_DIR) train_rgb, train_ela, train_labels = preload_images(splits['train']) val_rgb, val_ela, val_labels = preload_images(splits['val']) train_ds = make_dataset(train_rgb, train_ela, train_labels, shuffle=True) val_ds = make_dataset(val_rgb, val_ela, val_labels, shuffle=False) model = build_model() model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy']) steps_per_epoch = max(1, int(np.ceil(len(train_labels) / BATCH_SIZE))) validation_steps = max(1, int(np.ceil(len(val_labels) / BATCH_SIZE))) model.fit( train_ds, validation_data=val_ds, epochs=EPOCHS, steps_per_epoch=steps_per_epoch, validation_steps=validation_steps, verbose=1, ) model.save('M3_best.keras') print("Model saved as M3_best.keras")