utils/image_hasher.py

#!/usr/bin/env python3
"""
image_hasher.py

Provides robust image hashing utilities for MorphGuard to use in metrics collection
and content tracking. Supports multiple hashing algorithms for different use cases.
"""

import os
import io
import hashlib
import numpy as np
from PIL import Image
from typing import Union, Tuple, Optional, List, Dict, Any

# Try to import optional dependencies
try:
    import cv2
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False

try:
    import imagehash
    IMAGEHASH_AVAILABLE = True
except ImportError:
    IMAGEHASH_AVAILABLE = False


class ImageHasher:
    """Image hashing utility for MorphGuard"""
    
    def __init__(self, hash_size: int = 16, hash_algorithm: str = "phash"):
        """Initialize the image hasher
        
        Args:
            hash_size: Size of the hash (for perceptual hashing algorithms)
            hash_algorithm: Hashing algorithm to use 
                            ('md5', 'sha256', 'phash', 'dhash', 'ahash', 'whash')
        """
        self.hash_size = hash_size
        self.hash_algorithm = hash_algorithm.lower()
        
        # Check if required libraries are available for perceptual hashing
        if self.hash_algorithm in ('phash', 'dhash', 'ahash', 'whash'):
            if not IMAGEHASH_AVAILABLE:
                print(f"Warning: {hash_algorithm} requires the 'imagehash' package")
                print("Falling back to SHA-256 cryptographic hashing")
                self.hash_algorithm = 'sha256'
    
    def hash_image(self, image_path_or_data: Union[str, bytes, np.ndarray, Image.Image]) -> str:
        """Generate hash for an image
        
        Args:
            image_path_or_data: Image to hash (path, bytes, array, or PIL Image)
            
        Returns:
            String representation of the hash
        """
        # Load the image if needed
        img = self._load_image(image_path_or_data)
        
        if img is None:
            # Return a placeholder hash if image couldn't be loaded
            return hashlib.sha256(b'error_loading_image').hexdigest()
        
        # Generate hash based on algorithm
        if self.hash_algorithm == 'md5':
            return self._cryptographic_hash(img, 'md5')
        elif self.hash_algorithm == 'sha256':
            return self._cryptographic_hash(img, 'sha256')
        elif IMAGEHASH_AVAILABLE:
            # Use perceptual hashing if available
            if self.hash_algorithm == 'phash':
                img_hash = imagehash.phash(img, hash_size=self.hash_size)
            elif self.hash_algorithm == 'dhash':
                img_hash = imagehash.dhash(img, hash_size=self.hash_size)
            elif self.hash_algorithm == 'ahash':
                img_hash = imagehash.average_hash(img, hash_size=self.hash_size)
            elif self.hash_algorithm == 'whash':
                img_hash = imagehash.whash(img, hash_size=self.hash_size)
            else:
                # Default to perceptual hash
                img_hash = imagehash.phash(img, hash_size=self.hash_size)
            
            return str(img_hash)
        else:
            # Fall back to SHA-256 if perceptual hashing isn't available
            return self._cryptographic_hash(img, 'sha256')
    
    def hash_image_batch(self, images: List[Union[str, bytes, np.ndarray, Image.Image]]) -> List[str]:
        """Generate hashes for a batch of images
        
        Args:
            images: List of images to hash
            
        Returns:
            List of hash strings
        """
        return [self.hash_image(img) for img in images]
    
    def compare_images(self, 
                       image1: Union[str, bytes, np.ndarray, Image.Image],
                       image2: Union[str, bytes, np.ndarray, Image.Image]) -> float:
        """Compare two images and return similarity score
        
        Args:
            image1: First image to compare
            image2: Second image to compare
            
        Returns:
            Similarity score (0-1), where 1 is identical
        """
        # For cryptographic hashes, we can only do binary comparison
        if self.hash_algorithm in ('md5', 'sha256'):
            hash1 = self.hash_image(image1)
            hash2 = self.hash_image(image2)
            return 1.0 if hash1 == hash2 else 0.0
        
        # For perceptual hashes, we can calculate distance-based similarity
        if IMAGEHASH_AVAILABLE:
            img1 = self._load_image(image1)
            img2 = self._load_image(image2)
            
            if img1 is None or img2 is None:
                return 0.0
            
            if self.hash_algorithm == 'phash':
                hash1 = imagehash.phash(img1, hash_size=self.hash_size)
                hash2 = imagehash.phash(img2, hash_size=self.hash_size)
            elif self.hash_algorithm == 'dhash':
                hash1 = imagehash.dhash(img1, hash_size=self.hash_size)
                hash2 = imagehash.dhash(img2, hash_size=self.hash_size)
            elif self.hash_algorithm == 'ahash':
                hash1 = imagehash.average_hash(img1, hash_size=self.hash_size)
                hash2 = imagehash.average_hash(img2, hash_size=self.hash_size)
            elif self.hash_algorithm == 'whash':
                hash1 = imagehash.whash(img1, hash_size=self.hash_size)
                hash2 = imagehash.whash(img2, hash_size=self.hash_size)
            else:
                hash1 = imagehash.phash(img1, hash_size=self.hash_size)
                hash2 = imagehash.phash(img2, hash_size=self.hash_size)
            
            # Calculate normalized hamming distance-based similarity
            max_bits = self.hash_size * self.hash_size
            hamming_distance = hash1 - hash2
            similarity = 1.0 - (hamming_distance / max_bits)
            
            return float(similarity)
        
        # Fallback to direct image comparison if perceptual hashing isn't available
        return self._direct_image_comparison(image1, image2)
    
    def _load_image(self, image_path_or_data: Union[str, bytes, np.ndarray, Image.Image]) -> Optional[Image.Image]:
        """Load image from various input formats
        
        Args:
            image_path_or_data: Image to load (path, bytes, array, or PIL Image)
            
        Returns:
            PIL Image or None if loading failed
        """
        try:
            # Handle different input types
            if isinstance(image_path_or_data, str):
                # Path to image file
                return Image.open(image_path_or_data)
            elif isinstance(image_path_or_data, bytes):
                # Raw image bytes
                return Image.open(io.BytesIO(image_path_or_data))
            elif isinstance(image_path_or_data, np.ndarray):
                # NumPy array
                return Image.fromarray(image_path_or_data)
            elif isinstance(image_path_or_data, Image.Image):
                # Already a PIL Image
                return image_path_or_data
            else:
                print(f"Warning: Unsupported image type: {type(image_path_or_data)}")
                return None
        except Exception as e:
            print(f"Error loading image: {e}")
            return None
    
    def _cryptographic_hash(self, img: Image.Image, algorithm: str) -> str:
        """Generate a cryptographic hash of an image
        
        Args:
            img: PIL Image to hash
            algorithm: Hashing algorithm ('md5' or 'sha256')
            
        Returns:
            Hash string
        """
        # Convert to bytes for consistent hashing
        img_bytes = io.BytesIO()
        img.save(img_bytes, format='PNG')
        img_data = img_bytes.getvalue()
        
        # Apply hash function
        if algorithm == 'md5':
            return hashlib.md5(img_data).hexdigest()
        else:  # default to sha256
            return hashlib.sha256(img_data).hexdigest()
    
    def _direct_image_comparison(self, 
                                image1: Union[str, bytes, np.ndarray, Image.Image],
                                image2: Union[str, bytes, np.ndarray, Image.Image]) -> float:
        """Directly compare two images using pixel-wise comparison
        
        Args:
            image1: First image to compare
            image2: Second image to compare
            
        Returns:
            Similarity score (0-1)
        """
        img1 = self._load_image(image1)
        img2 = self._load_image(image2)
        
        if img1 is None or img2 is None:
            return 0.0
        
        # Resize images to the same dimensions
        size = (128, 128)  # Small size for faster comparison
        img1 = img1.resize(size, Image.LANCZOS)
        img2 = img2.resize(size, Image.LANCZOS)
        
        # Convert to grayscale for simplicity
        img1 = img1.convert('L')
        img2 = img2.convert('L')
        
        # Convert to numpy arrays
        arr1 = np.array(img1)
        arr2 = np.array(img2)
        
        # Calculate mean squared error
        mse = np.mean((arr1 - arr2) ** 2)
        if mse == 0:
            return 1.0
        
        # Convert MSE to similarity score (0-1)
        max_mse = 255.0 ** 2  # Maximum possible MSE
        similarity = 1.0 - (mse / max_mse)
        
        return float(similarity)


def hash_file(file_path: str, algorithm: str = 'sha256') -> str:
    """Generate a hash for any file
    
    Args:
        file_path: Path to the file
        algorithm: Hashing algorithm ('md5' or 'sha256')
        
    Returns:
        Hash string
    """
    if not os.path.exists(file_path):
        return ""
    
    try:
        with open(file_path, 'rb') as f:
            file_data = f.read()
            
        if algorithm == 'md5':
            return hashlib.md5(file_data).hexdigest()
        else:  # default to sha256
            return hashlib.sha256(file_data).hexdigest()
    except Exception as e:
        print(f"Error hashing file: {e}")
        return ""


# Convenience functions with default settings
def get_image_hash(image_path_or_data: Union[str, bytes, np.ndarray, Image.Image], 
                  algorithm: str = 'phash') -> str:
    """Get a hash for an image using the specified algorithm
    
    Args:
        image_path_or_data: Image to hash (path, bytes, array, or PIL Image)
        algorithm: Hashing algorithm ('md5', 'sha256', 'phash', 'dhash', 'ahash', 'whash')
        
    Returns:
        Hash string
    """
    hasher = ImageHasher(hash_algorithm=algorithm)
    return hasher.hash_image(image_path_or_data)


def compare_images(image1: Union[str, bytes, np.ndarray, Image.Image],
                 image2: Union[str, bytes, np.ndarray, Image.Image],
                 algorithm: str = 'phash') -> float:
    """Compare two images and return similarity score
    
    Args:
        image1: First image to compare
        image2: Second image to compare
        algorithm: Hashing algorithm for comparison
        
    Returns:
        Similarity score (0-1)
    """
    hasher = ImageHasher(hash_algorithm=algorithm)
    return hasher.compare_images(image1, image2)


# Function to get a comprehensive set of hashes for an image
def get_all_hashes(image_path_or_data: Union[str, bytes, np.ndarray, Image.Image]) -> Dict[str, str]:
    """Get multiple hashes for a single image
    
    Args:
        image_path_or_data: Image to hash
        
    Returns:
        Dictionary of hash algorithm -> hash value
    """
    result = {}
    
    # Cryptographic hashes
    md5_hasher = ImageHasher(hash_algorithm='md5')
    result['md5'] = md5_hasher.hash_image(image_path_or_data)
    
    sha256_hasher = ImageHasher(hash_algorithm='sha256')
    result['sha256'] = sha256_hasher.hash_image(image_path_or_data)
    
    # Perceptual hashes if available
    if IMAGEHASH_AVAILABLE:
        phash_hasher = ImageHasher(hash_algorithm='phash')
        result['phash'] = phash_hasher.hash_image(image_path_or_data)
        
        dhash_hasher = ImageHasher(hash_algorithm='dhash')
        result['dhash'] = dhash_hasher.hash_image(image_path_or_data)
        
        ahash_hasher = ImageHasher(hash_algorithm='ahash')
        result['ahash'] = ahash_hasher.hash_image(image_path_or_data)
        
        whash_hasher = ImageHasher(hash_algorithm='whash')
        result['whash'] = whash_hasher.hash_image(image_path_or_data)
    
    return result