commented version of PR: "Re-sizable preview window & --viewer_size CLI argument #175" https://github.com/HyperGAN/HyperGAN/pull/175

pygame_viewer_commented.py

"""
Opens a window that displays an image.
Usage:

    from viewer import GlobalViewer
    GlobalViewer.update(image)

"""
import numpy as np



class PygameViewer:

    def __init__(self, title="HyperGAN", viewer_size=1, enabled=True):
        self.screen = None
        self.title = title
        self.viewer_size = viewer_size
        self.enabled = enabled


    def update(self, image):
        if not self.enabled: return

        image = np.transpose(image, [1, 0,2])

        if not self.screen:
            import pygame
            self.pg = pygame
            #set viewer size limit to stop pygame errors
            if self.viewer_size <= 0:
                self.viewer_size = 0.1
            #multiply base size by size arg
            self.size = [int(image.shape[0] * self.viewer_size), int(image.shape[1] * self.viewer_size)]
            #calculate aspect ratios for w:h and h:w
            self.aspect_w = image.shape[1] / image.shape[0]
            self.aspect_h = image.shape[0] / image.shape[1]
            self.temp_size = self.size
            self.screen = self.pg.display.set_mode(self.size,self.pg.RESIZABLE)
            self.pg.display.set_caption(self.title)

        #loop through all events that happened since last update
        for event in self.pg.event.get():
            #if we resized
            if event.type == self.pg.VIDEORESIZE:
                #if it was changed horizontally
                if self.size[0] != event.size[0]:
                    #calc new size with respect for aspect ratio
                    self.temp_size = [event.size[0], int(event.size[0] * self.aspect_w)]
                #if it was changed vertically
                elif self.size[1] != event.size[1]:
                    #calc new size with respect for aspect ratio
                    self.temp_size = [int(event.size[1] * self.aspect_h), event.size[1]]
            #when the user stops resizing (pygame doesnt respond to .set_mode() while resizing so we do it after)
            elif event.type == self.pg.ACTIVEEVENT and event.state == 2 and event.gain == 1:
                #set window size
                self.size = self.temp_size
                self.screen = self.pg.display.set_mode(self.size,self.pg.RESIZABLE)   
                
        surface = self.pg.Surface([image.shape[0],image.shape[1]])
        self.pg.surfarray.blit_array(surface, image)
        self.screen.blit(self.pg.transform.scale(surface,self.size),(0,0))
        self.pg.display.flip()