| 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 |
|
|
| |
| 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") |
|
|
