main.py

# main.py

import os
import cv2
import yaml
import logging
from dip_features import compute_sharpness, compute_edge_density, extract_sift_features, compute_tampered_ratio_and_bbox, detect_objects
from csv_export import save_features_to_csv

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def load_config(config_path="config.yaml"):
    try:
        with open(config_path, "r") as f:
            config = yaml.safe_load(f)
        logger.info("Configuration loaded from %s.", config_path)
        return config
    except Exception as e:
        logger.error("Failed to load configuration: %s", e)
        raise

def process_image(image_path, config):
    """
    Load an image, preprocess it, and extract DIP features.
    Returns a dictionary of features and a bounding box.
    """
    try:
        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if image is None:
            logger.error("Failed to load image: %s", image_path)
            return None, (0, 0, 0, 0)
        size = config.get("image_size", 500)
        image = cv2.resize(image, (size, size))
        
        features = {
            'sharpness': compute_sharpness(image),
            'edge_density': compute_edge_density(image, config.get("canny_threshold1", 100), config.get("canny_threshold2", 200)),
            'sift_count': extract_sift_features(image),
            'object_count': detect_objects(image)
        }
        tampered_ratio, bbox = compute_tampered_ratio_and_bbox(image, config.get("block_size", 64),
                                                               config.get("tampered_ssim_threshold", 0.85))
        features['tampered_ratio'] = tampered_ratio
        
        return features, bbox
    except Exception as e:
        logger.error("Error processing image %s: %s", image_path, e)
        return None, (0, 0, 0, 0)

def process_dataset(config):
    dataset_dir = config.get("dataset_dir", "./dataset/CASIA2")
    output_csv = config.get("output_csv", "forgery_features.csv")
    
    for label in ['authentic', 'forged']:
        folder = os.path.join(dataset_dir, label)
        if not os.path.isdir(folder):
            logger.warning("Directory not found: %s", folder)
            continue
        for root, _, files in os.walk(folder):
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')):
                    image_path = os.path.join(root, file)
                    features, bbox = process_image(image_path, config)
                    if features:
                        save_features_to_csv(file, features, bbox, label, output_csv)
                        logger.info("Processed %s", file)
    logger.info("Feature extraction complete. CSV saved at: %s", output_csv)

if __name__ == "__main__":
    config = load_config("config.yaml")
    process_dataset(config)