PHiLIP/display_image.py at main · freddysongg/PHiLIP · GitHub

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
# display_image.py

import matplotlib.pyplot as plt
import numpy as np
import os
import logging
from typing import List, Optional
from config import IMAGE_FOLDER, LOG_FORMAT, LOG_DATE_FORMAT

logging.basicConfig(level=logging.INFO, format=LOG_FORMAT, datefmt=LOG_DATE_FORMAT)
logger = logging.getLogger(__name__)

def display_and_select_image(images: List[np.ndarray], resolution: int, iteration: int) -> Optional[List[np.ndarray]]:
    """
    Displays images, prompts the user to select their favorite, and saves all images.

    Args:
        images: List of images to display and save.
        resolution: Resolution of the images.
        iteration: Current iteration number.

    Returns:
        Optional[List[np.ndarray]]: List of selected images or None if no selection is made.
    """
    num_images = len(images)
    fig, axs = plt.subplots(1, num_images, figsize=(5 * num_images, 5))
    axs = axs if num_images > 1 else [axs]

    for i, img in enumerate(images):
        axs[i].imshow(img)
        axs[i].axis('off')
        axs[i].set_title(f"Image {i+1}")

    plt.tight_layout()
    plt.show()

    save_images(images, resolution)

    return get_user_selection(images, num_images, resolution)

def save_images(images: List[np.ndarray], resolution: int, final: bool = False) -> None:
    """
    Save all generated images.

    Args:
        images: List of images to save.
        resolution: Resolution of the images.
        final: Whether this is the final enhanced image.
    """
    num_images = len(images)
    logger.info(f"Saving {num_images} image{'s' if num_images > 1 else ''} at {resolution}x{resolution} resolution")

    for i, img in enumerate(images):
        file_name = f"final-enhanced-{resolution}.png" if final else f"{resolution}-{i+1}.png"
        file_path = os.path.join(IMAGE_FOLDER, file_name)
        try:
            plt.imsave(file_path, img)
            logger.info(f"Image {file_name} saved successfully")
        except Exception as e:
            logger.error(f"Failed to save image {file_name}: {str(e)}")

    if final:
        logger.info(f"Final enhanced image saved as {file_name}")
    else:
        logger.info(f"All {num_images} images saved in {IMAGE_FOLDER}")

def get_user_selection(images: List[np.ndarray], num_images: int, resolution: int) -> Optional[List[np.ndarray]]:
    """
    Get user's image selection.

    Args:
        images: List of images to choose from.
        num_images: Total number of images.
        resolution: Resolution of the images.

    Returns:
        Optional[List[np.ndarray]]: List of selected images or None if no selection is made.
    """
    while True:
        choice = input(f"Select your favorite images (1-{num_images}), separated by commas, or type 'stop' to exit: ").strip()

        if choice.lower() == 'stop':
            return None

        try:
            selected_indices = [int(x) - 1 for x in choice.split(',')]
            selected_images = [images[i] for i in selected_indices if 0 <= i < num_images]

            if not selected_images:
                raise ValueError("No valid images selected")

            rename_selected_images(selected_indices, resolution)
            return selected_images
        except ValueError as e:
            logger.warning(f"Invalid input: {e}")
            print(f"Invalid input: {e}. Please enter numbers between 1 and {num_images} separated by commas or 'stop' to exit.")

def rename_selected_images(selected_indices: List[int], resolution: int) -> None:
    """
    Rename selected images.

    Args:
        selected_indices: Indices of selected images.
        resolution: Resolution of the images.
    """
    for i in selected_indices:
        old_path = os.path.join(IMAGE_FOLDER, f"{resolution}-{i+1}.png")
        new_path = os.path.join(IMAGE_FOLDER, f"{resolution}-selected-{i+1}.png")
        try:
            os.rename(old_path, new_path)
            logger.info(f"Image {i+1} renamed to '{os.path.basename(new_path)}'")
        except Exception as e:
            logger.error(f"Failed to rename image {i+1}: {str(e)}")