File size: 8,602 Bytes
b56342d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
import numpy as np
import os
import time
from . import util
from . import html
# from pdb import set_trace as st
import matplotlib.pyplot as plt
import math
# from IPython import embed

def zoom_to_res(img,res=256,order=0,axis=0):
    # img   3xXxX
    from scipy.ndimage import zoom
    zoom_factor = res/img.shape[1]
    if(axis==0):
        return zoom(img,[1,zoom_factor,zoom_factor],order=order)
    elif(axis==2):
        return zoom(img,[zoom_factor,zoom_factor,1],order=order)

class Visualizer():
    def __init__(self, opt):
        # self.opt = opt
        self.display_id = opt.display_id
        # self.use_html = opt.is_train and not opt.no_html
        self.win_size = opt.display_winsize
        self.name = opt.name
        self.display_cnt = 0 # display_current_results counter
        self.display_cnt_high = 0
        self.use_html = opt.use_html

        if self.display_id > 0:
            import visdom
            self.vis = visdom.Visdom(port = opt.display_port)

        self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
        util.mkdirs([self.web_dir,])
        if self.use_html:
            self.img_dir = os.path.join(self.web_dir, 'images')
            print('create web directory %s...' % self.web_dir)
            util.mkdirs([self.img_dir,])

    # |visuals|: dictionary of images to display or save
    def display_current_results(self, visuals, epoch, nrows=None, res=256):
        if self.display_id > 0: # show images in the browser
            title = self.name
            if(nrows is None):
                nrows = int(math.ceil(len(visuals.items()) / 2.0))
            images = []
            idx = 0
            for label, image_numpy in visuals.items():
                title += " | " if idx % nrows == 0 else ", "
                title += label
                img = image_numpy.transpose([2, 0, 1])
                img = zoom_to_res(img,res=res,order=0)
                images.append(img)
                idx += 1
            if len(visuals.items()) % 2 != 0:
                white_image = np.ones_like(image_numpy.transpose([2, 0, 1]))*255
                white_image = zoom_to_res(white_image,res=res,order=0)
                images.append(white_image)
            self.vis.images(images, nrow=nrows, win=self.display_id + 1,
                            opts=dict(title=title))

        if self.use_html: # save images to a html file
            for label, image_numpy in visuals.items():
                img_path = os.path.join(self.img_dir, 'epoch%.3d_cnt%.6d_%s.png' % (epoch, self.display_cnt, label))
                util.save_image(zoom_to_res(image_numpy, res=res, axis=2), img_path)

            self.display_cnt += 1
            self.display_cnt_high = np.maximum(self.display_cnt_high, self.display_cnt)

            # update website
            webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, reflesh=1)
            for n in range(epoch, 0, -1):
                webpage.add_header('epoch [%d]' % n)
                if(n==epoch):
                    high = self.display_cnt
                else:
                    high = self.display_cnt_high
                for c in range(high-1,-1,-1):
                    ims = []
                    txts = []
                    links = []

                    for label, image_numpy in visuals.items():
                        img_path = 'epoch%.3d_cnt%.6d_%s.png' % (n, c, label)
                        ims.append(os.path.join('images',img_path))
                        txts.append(label)
                        links.append(os.path.join('images',img_path))
                    webpage.add_images(ims, txts, links, width=self.win_size)
            webpage.save()

    # save errors into a directory
    def plot_current_errors_save(self, epoch, counter_ratio, opt, errors,keys='+ALL',name='loss', to_plot=False):
        if not hasattr(self, 'plot_data'):
            self.plot_data = {'X':[],'Y':[], 'legend':list(errors.keys())}
        self.plot_data['X'].append(epoch + counter_ratio)
        self.plot_data['Y'].append([errors[k] for k in self.plot_data['legend']])

        # embed()
        if(keys=='+ALL'):
            plot_keys = self.plot_data['legend']
        else:
            plot_keys = keys

        if(to_plot):
            (f,ax) = plt.subplots(1,1)
        for (k,kname) in enumerate(plot_keys):
            kk = np.where(np.array(self.plot_data['legend'])==kname)[0][0]
            x = self.plot_data['X']
            y = np.array(self.plot_data['Y'])[:,kk]
            if(to_plot):
                ax.plot(x, y, 'o-', label=kname)
            np.save(os.path.join(self.web_dir,'%s_x')%kname,x)
            np.save(os.path.join(self.web_dir,'%s_y')%kname,y)

        if(to_plot):
            plt.legend(loc=0,fontsize='small')
            plt.xlabel('epoch')
            plt.ylabel('Value')
            f.savefig(os.path.join(self.web_dir,'%s.png'%name))
            f.clf()
            plt.close()

    # errors: dictionary of error labels and values
    def plot_current_errors(self, epoch, counter_ratio, opt, errors):
        if not hasattr(self, 'plot_data'):
            self.plot_data = {'X':[],'Y':[], 'legend':list(errors.keys())}
        self.plot_data['X'].append(epoch + counter_ratio)
        self.plot_data['Y'].append([errors[k] for k in self.plot_data['legend']])
        self.vis.line(
            X=np.stack([np.array(self.plot_data['X'])]*len(self.plot_data['legend']),1),
            Y=np.array(self.plot_data['Y']),
            opts={
                'title': self.name + ' loss over time',
                'legend': self.plot_data['legend'],
                'xlabel': 'epoch',
                'ylabel': 'loss'},
            win=self.display_id)

    # errors: same format as |errors| of plotCurrentErrors
    def print_current_errors(self, epoch, i, errors, t, t2=-1, t2o=-1, fid=None):
        message = '(ep: %d, it: %d, t: %.3f[s], ept: %.2f/%.2f[h]) ' % (epoch, i, t, t2o, t2)
        message += (', ').join(['%s: %.3f' % (k, v) for k, v in errors.items()])

        print(message)
        if(fid is not None):
            fid.write('%s\n'%message)


    # save image to the disk
    def save_images_simple(self, webpage, images, names, in_txts, prefix='', res=256):
        image_dir = webpage.get_image_dir()
        ims = []
        txts = []
        links = []

        for name, image_numpy, txt in zip(names, images, in_txts):
            image_name = '%s_%s.png' % (prefix, name)
            save_path = os.path.join(image_dir, image_name)
            if(res is not None):
                util.save_image(zoom_to_res(image_numpy,res=res,axis=2), save_path)
            else:
                util.save_image(image_numpy, save_path)

            ims.append(os.path.join(webpage.img_subdir,image_name))
            # txts.append(name)
            txts.append(txt)
            links.append(os.path.join(webpage.img_subdir,image_name))
        # embed()
        webpage.add_images(ims, txts, links, width=self.win_size)

    # save image to the disk
    def save_images(self, webpage, images, names, image_path, title=''):
        image_dir = webpage.get_image_dir()
        # short_path = ntpath.basename(image_path)
        # name = os.path.splitext(short_path)[0]
        # name = short_path
        # webpage.add_header('%s, %s' % (name, title))
        ims = []
        txts = []
        links = []

        for label, image_numpy in zip(names, images):
            image_name = '%s.jpg' % (label,)
            save_path = os.path.join(image_dir, image_name)
            util.save_image(image_numpy, save_path)

            ims.append(image_name)
            txts.append(label)
            links.append(image_name)
        webpage.add_images(ims, txts, links, width=self.win_size)

    # save image to the disk
    # def save_images(self, webpage, visuals, image_path, short=False):
    #     image_dir = webpage.get_image_dir()
    #     if short:
    #         short_path = ntpath.basename(image_path)
    #         name = os.path.splitext(short_path)[0]
    #     else:
    #         name = image_path

    #     webpage.add_header(name)
    #     ims = []
    #     txts = []
    #     links = []

    #     for label, image_numpy in visuals.items():
    #         image_name = '%s_%s.png' % (name, label)
    #         save_path = os.path.join(image_dir, image_name)
    #         util.save_image(image_numpy, save_path)

    #         ims.append(image_name)
    #         txts.append(label)
    #         links.append(image_name)
    #     webpage.add_images(ims, txts, links, width=self.win_size)