Upload utils.py

utils.py

@@ -0,0 +1,391 @@
+"""
+Utility functions for the Interior Style Transfer Pipeline
+"""
+import cv2
+import numpy as np
+from PIL import Image
+import os
+from typing import Tuple, List, Optional, Union
+import json
+from pathlib import Path
+
+def load_image_safe(image_path: str, target_size: Tuple[int, int] = None) -> np.ndarray:
+    """
+    Safely load an image with error handling
+    
+    Args:
+        image_path: Path to the image file
+        target_size: Optional target size (width, height)
+        
+    Returns:
+        Loaded image as numpy array
+        
+    Raises:
+        ValueError: If image cannot be loaded
+    """
+    if not os.path.exists(image_path):
+        raise ValueError(f"Image file not found: {image_path}")
+    
+    # Try to load with OpenCV first
+    image = cv2.imread(image_path)
+    if image is None:
+        # Fallback to PIL
+        try:
+            pil_image = Image.open(image_path)
+            image = np.array(pil_image)
+            if len(image.shape) == 3 and image.shape[2] == 3:
+                image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+            elif len(image.shape) == 3 and image.shape[2] == 4:
+                image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)
+        except Exception as e:
+            raise ValueError(f"Could not load image {image_path}: {e}")
+    
+    if target_size:
+        image = cv2.resize(image, target_size)
+    
+    return image
+
+def save_image_safe(image: np.ndarray, output_path: str, 
+                   quality: int = 95) -> bool:
+    """
+    Safely save an image with error handling
+    
+    Args:
+        image: Image to save as numpy array
+        output_path: Output file path
+        quality: JPEG quality (1-100)
+        
+    Returns:
+        True if successful, False otherwise
+    """
+    try:
+        # Ensure output directory exists
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        
+        # Save with OpenCV
+        success = cv2.imwrite(output_path, image)
+        
+        if not success:
+            # Fallback to PIL
+            if len(image.shape) == 3 and image.shape[2] == 3:
+                pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+            else:
+                pil_image = Image.fromarray(image)
+            
+            pil_image.save(output_path, quality=quality)
+            success = True
+        
+        return success
+    except Exception as e:
+        print(f"Error saving image to {output_path}: {e}")
+        return False
+
+def validate_image_pair(user_room: np.ndarray, inspiration_room: np.ndarray) -> Tuple[bool, str]:
+    """
+    Validate that two images are suitable for style transfer
+    
+    Args:
+        user_room: User room image
+        inspiration_room: Inspiration room image
+        
+    Returns:
+        Tuple of (is_valid, error_message)
+    """
+    # Check image dimensions
+    if user_room.shape != inspiration_room.shape:
+        return False, f"Image dimensions don't match: {user_room.shape} vs {inspiration_room.shape}"
+    
+    # Check minimum size
+    min_size = 256
+    if user_room.shape[0] < min_size or user_room.shape[1] < min_size:
+        return False, f"Images too small. Minimum size: {min_size}x{min_size}"
+    
+    # Check aspect ratio (should be roughly square for best results)
+    aspect_ratio = user_room.shape[1] / user_room.shape[0]
+    if aspect_ratio < 0.5 or aspect_ratio > 2.0:
+        return False, f"Extreme aspect ratio: {aspect_ratio:.2f}. Square images work best."
+    
+    # Check if images are too dark or too bright
+    user_brightness = np.mean(cv2.cvtColor(user_room, cv2.COLOR_BGR2GRAY))
+    inspiration_brightness = np.mean(cv2.cvtColor(inspiration_room, cv2.COLOR_BGR2GRAY))
+    
+    if user_brightness < 30 or user_brightness > 225:
+        return False, f"User room too {'dark' if user_brightness < 30 else 'bright'}: {user_brightness:.1f}"
+    
+    if inspiration_brightness < 30 or inspiration_brightness > 225:
+        return False, f"Inspiration room too {'dark' if inspiration_brightness < 30 else 'bright'}: {inspiration_brightness:.1f}"
+    
+    return True, "Images are valid for style transfer"
+
+def create_comparison_image(original: np.ndarray, result: np.ndarray, 
+                           title: str = "Style Transfer Comparison") -> np.ndarray:
+    """
+    Create a side-by-side comparison image
+    
+    Args:
+        original: Original user room image
+        result: Style transfer result
+        title: Title for the comparison
+        
+    Returns:
+        Comparison image
+    """
+    # Ensure both images have the same dimensions
+    if original.shape != result.shape:
+        result = cv2.resize(result, (original.shape[1], original.shape[0]))
+    
+    # Create comparison image
+    comparison = np.hstack([original, result])
+    
+    # Add title
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    font_scale = 1.0
+    thickness = 2
+    
+    # Calculate text position
+    text_size = cv2.getTextSize(title, font, font_scale, thickness)[0]
+    text_x = (comparison.shape[1] - text_size[0]) // 2
+    text_y = 50
+    
+    # Add background for text
+    cv2.rectangle(comparison, (text_x - 10, text_y - 30), 
+                  (text_x + text_size[0] + 10, text_y + 10), (255, 255, 255), -1)
+    
+    # Add text
+    cv2.putText(comparison, title, (text_x, text_y), font, font_scale, (0, 0, 0), thickness)
+    
+    # Add labels
+    cv2.putText(comparison, "Original", (50, comparison.shape[0] - 30), 
+                font, 0.7, (255, 255, 255), 2)
+    cv2.putText(comparison, "Result", (original.shape[1] + 50, comparison.shape[0] - 30), 
+                font, 0.7, (255, 255, 255), 2)
+    
+    return comparison
+
+def enhance_image_quality(image: np.ndarray, 
+                         sharpness: float = 0.3,
+                         contrast: float = 1.1,
+                         saturation: float = 1.1) -> np.ndarray:
+    """
+    Enhance image quality with various filters
+    
+    Args:
+        image: Input image
+        sharpness: Sharpening strength (0.0 to 1.0)
+        contrast: Contrast multiplier
+        saturation: Saturation multiplier
+        
+    Returns:
+        Enhanced image
+    """
+    enhanced = image.copy()
+    
+    # Sharpening
+    if sharpness > 0:
+        kernel = np.array([[-1, -1, -1],
+                          [-1,  9, -1],
+                          [-1, -1, -1]]) * sharpness
+        enhanced = cv2.filter2D(enhanced, -1, kernel)
+    
+    # Contrast adjustment
+    if contrast != 1.0:
+        enhanced = np.clip(enhanced * contrast, 0, 255).astype(np.uint8)
+    
+    # Saturation adjustment
+    if saturation != 1.0:
+        hsv = cv2.cvtColor(enhanced, cv2.COLOR_BGR2HSV).astype(np.float32)
+        hsv[:, :, 1] = np.clip(hsv[:, :, 1] * saturation, 0, 255)
+        enhanced = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
+    
+    return enhanced
+
+def create_progress_bar(total: int, description: str = "Processing") -> callable:
+    """
+    Create a simple progress bar function
+    
+    Args:
+        total: Total number of steps
+        description: Description of the process
+        
+    Returns:
+        Function to update progress
+    """
+    def update_progress(current: int):
+        percentage = (current / total) * 100
+        bar_length = 30
+        filled_length = int(bar_length * current // total)
+        bar = '█' * filled_length + '-' * (bar_length - filled_length)
+        print(f'\r{description}: |{bar}| {percentage:.1f}% ({current}/{total})', end='')
+        if current == total:
+            print()
+    
+    return update_progress
+
+def save_metadata(metadata: dict, output_path: str) -> bool:
+    """
+    Save metadata to JSON file
+    
+    Args:
+        metadata: Dictionary of metadata
+        output_path: Output file path
+        
+    Returns:
+        True if successful, False otherwise
+    """
+    try:
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        
+        with open(output_path, 'w') as f:
+            json.dump(metadata, f, indent=2, default=str)
+        
+        return True
+    except Exception as e:
+        print(f"Error saving metadata to {output_path}: {e}")
+        return False
+
+def load_metadata(metadata_path: str) -> Optional[dict]:
+    """
+    Load metadata from JSON file
+    
+    Args:
+        metadata_path: Path to metadata file
+        
+    Returns:
+        Loaded metadata dictionary or None if failed
+    """
+    try:
+        with open(metadata_path, 'r') as f:
+            return json.load(f)
+    except Exception as e:
+        print(f"Error loading metadata from {metadata_path}: {e}")
+        return None
+
+def calculate_image_similarity(img1: np.ndarray, img2: np.ndarray) -> float:
+    """
+    Calculate similarity between two images using structural similarity
+    
+    Args:
+        img1: First image
+        img2: Second image
+        
+    Returns:
+        Similarity score (0.0 to 1.0, higher is more similar)
+    """
+    try:
+        from skimage.metrics import structural_similarity as ssim
+        
+        # Ensure same dimensions
+        if img1.shape != img2.shape:
+            img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))
+        
+        # Convert to grayscale for SSIM
+        gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
+        gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
+        
+        # Calculate SSIM
+        similarity = ssim(gray1, gray2)
+        return max(0.0, similarity)  # Ensure non-negative
+        
+    except ImportError:
+        # Fallback to simple MSE-based similarity
+        if img1.shape != img2.shape:
+            img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))
+        
+        mse = np.mean((img1.astype(np.float32) - img2.astype(np.float32)) ** 2)
+        max_mse = 255 ** 2
+        similarity = 1.0 - (mse / max_mse)
+        return max(0.0, similarity)
+
+def create_thumbnail(image: np.ndarray, max_size: int = 200) -> np.ndarray:
+    """
+    Create a thumbnail version of an image
+    
+    Args:
+        image: Input image
+        max_size: Maximum dimension size
+        
+    Returns:
+        Thumbnail image
+    """
+    height, width = image.shape[:2]
+    
+    if height <= max_size and width <= max_size:
+        return image.copy()
+    
+    # Calculate new dimensions maintaining aspect ratio
+    if height > width:
+        new_height = max_size
+        new_width = int(width * max_size / height)
+    else:
+        new_width = max_size
+        new_height = int(height * max_size / width)
+    
+    thumbnail = cv2.resize(image, (new_width, new_height))
+    return thumbnail
+
+def batch_resize_images(images: List[np.ndarray], 
+                       target_size: Tuple[int, int]) -> List[np.ndarray]:
+    """
+    Resize a list of images to the same target size
+    
+    Args:
+        images: List of input images
+        target_size: Target size (width, height)
+        
+    Returns:
+        List of resized images
+    """
+    resized_images = []
+    
+    for image in images:
+        resized = cv2.resize(image, target_size)
+        resized_images.append(resized)
+    
+    return resized_images
+
+def create_image_grid(images: List[np.ndarray], 
+                     grid_size: Tuple[int, int] = None) -> np.ndarray:
+    """
+    Create a grid layout of images
+    
+    Args:
+        images: List of images to arrange in grid
+        grid_size: Grid dimensions (rows, cols). If None, auto-calculate
+        
+    Returns:
+        Grid image
+    """
+    if not images:
+        return np.array([])
+    
+    if grid_size is None:
+        # Auto-calculate grid size
+        n_images = len(images)
+        cols = int(np.ceil(np.sqrt(n_images)))
+        rows = int(np.ceil(n_images / cols))
+        grid_size = (rows, cols)
+    
+    rows, cols = grid_size
+    
+    # Ensure all images have the same size
+    target_size = (images[0].shape[1], images[0].shape[0])
+    resized_images = batch_resize_images(images, target_size)
+    
+    # Create grid
+    grid_rows = []
+    for i in range(rows):
+        row_images = []
+        for j in range(cols):
+            idx = i * cols + j
+            if idx < len(resized_images):
+                row_images.append(resized_images[idx])
+            else:
+                # Fill empty space with black
+                empty_image = np.zeros((target_size[1], target_size[0], 3), dtype=np.uint8)
+                row_images.append(empty_image)
+        
+        row = np.hstack(row_images)
+        grid_rows.append(row)
+    
+    grid = np.vstack(grid_rows)
+    return grid