Spaces:
Running
Running
File size: 3,403 Bytes
ad8cacf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
"""
Overview image generation functionality.
"""
import os
import numpy as np
import matplotlib.pyplot as plt
import cv2
from loguru import logger
def create_overview_image(refined_dir='refined_animals', overview_filename='overview.png'):
"""
Create an overview image showing all processed animal illustrations in a grid.
Args:
refined_dir (str): Directory containing the refined animal images
overview_filename (str): Filename for the overview image
Returns:
str: Path to the overview image
"""
# Get list of image files
image_files = [f for f in os.listdir(refined_dir)
if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
if not image_files:
logger.warning(f"No images found in {refined_dir}")
return None
try:
# Sort files for consistent display
image_files.sort()
# Calculate grid dimensions
n_images = len(image_files)
grid_size = int(np.ceil(np.sqrt(n_images)))
rows = cols = grid_size
# Create figure
plt.figure(figsize=(15, 15))
# Add each image to the grid
for i, filename in enumerate(image_files):
if i >= rows * cols:
break
# Read and process image
img_path = os.path.join(refined_dir, filename)
img = load_and_convert_image(img_path)
if img is not None:
# Add to plot
plt.subplot(rows, cols, i + 1)
plt.imshow(img)
plt.title(filename)
plt.axis('off')
# Adjust layout and save
plt.tight_layout()
overview_path = os.path.join(refined_dir, overview_filename)
plt.savefig(overview_path)
plt.close()
logger.info(f"Overview image saved to {overview_path}")
return overview_path
except Exception as e:
logger.error(f"Error creating overview image: {str(e)}")
return None
def load_and_convert_image(image_path):
"""
Load and convert image for matplotlib display.
Args:
image_path (str): Path to image file
Returns:
numpy.ndarray: RGB image array or None on failure
"""
try:
# Read image
img = cv2.imread(image_path)
if img is None:
logger.warning(f"Could not read image: {image_path}")
return None
# Convert from BGR to RGB
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
except Exception as e:
logger.error(f"Error loading image {image_path}: {str(e)}")
return None
def create_thumbnail(img, max_size=200):
"""
Create thumbnail version of image.
Args:
img (numpy.ndarray): Input image
max_size (int): Maximum dimension
Returns:
numpy.ndarray: Resized image
"""
# Get current dimensions
height, width = img.shape[:2]
# Calculate new dimensions
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))
# Resize image
return cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
|