first version

2022-07-16 15:46:17 +08:00
commit df69507ef2
32 changed files with 2336 additions and 0 deletions
@@ -0,0 +1,66 @@
+import dominate
+from dominate.tags import *
+import os
+
+
+class HTML:
+    def __init__(self, web_dir, title, image_subdir='', reflesh=0):
+        self.title = title
+        self.web_dir = web_dir
+        # self.img_dir = os.path.join(self.web_dir, )
+        self.img_subdir = image_subdir
+        self.img_dir = os.path.join(self.web_dir, image_subdir)
+        if not os.path.exists(self.web_dir):
+            os.makedirs(self.web_dir)
+        if not os.path.exists(self.img_dir):
+            os.makedirs(self.img_dir)
+        # print(self.img_dir)
+
+        self.doc = dominate.document(title=title)
+        if reflesh > 0:
+            with self.doc.head:
+                meta(http_equiv="reflesh", content=str(reflesh))
+
+    def get_image_dir(self):
+        return self.img_dir
+
+    def add_header(self, str):
+        with self.doc:
+            h3(str)
+
+    def add_table(self, border=1):
+        self.t = table(border=border, style="table-layout: fixed;")
+        self.doc.add(self.t)
+
+    def add_images(self, ims, txts, links, width=400):
+        self.add_table()
+        with self.t:
+            with tr():
+                for im, txt, link in zip(ims, txts, links):
+                    with td(style="word-wrap: break-word;", halign="center", valign="top"):
+                        with p():
+                            with a(href=os.path.join(link)):
+                                img(style="width:%dpx" % width, src=os.path.join(im))
+                            br()
+                            p(txt)
+
+    def save(self,file='index'):
+        html_file = '%s/%s.html' % (self.web_dir,file)
+        f = open(html_file, 'wt')
+        f.write(self.doc.render())
+        f.close()
+
+
+if __name__ == '__main__':
+    html = HTML('web/', 'test_html')
+    html.add_header('hello world')
+
+    ims = []
+    txts = []
+    links = []
+    for n in range(4):
+        ims.append('image_%d.png' % n)
+        txts.append('text_%d' % n)
+        links.append('image_%d.png' % n)
+    html.add_images(ims, txts, links)
+    html.save()
@@ -0,0 +1,33 @@
+import random
+import numpy as np
+import torch
+from torch.autograd import Variable
+from collections import deque
+
+
+class ImagePool():
+    def __init__(self, pool_size):
+        self.pool_size = pool_size
+        self.sample_size = pool_size
+        if self.pool_size > 0:
+            self.num_imgs = 0
+            self.images = deque()
+
+    def add(self, images):
+        if self.pool_size == 0:
+            return images
+        for image in images.data:
+            image = torch.unsqueeze(image, 0)
+            if self.num_imgs < self.pool_size:
+                self.num_imgs = self.num_imgs + 1
+                self.images.append(image)
+            else:
+                self.images.popleft()
+                self.images.append(image)
+
+    def query(self):
+        if len(self.images) > self.sample_size:
+            return_images = list(random.sample(self.images, self.sample_size))
+        else:
+            return_images = list(self.images)
+        return torch.cat(return_images, 0)
@@ -0,0 +1,54 @@
+import math
+from math import exp
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
+    return gauss / gauss.sum()
+
+
+def create_window(window_size, channel):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+
+def SSIM(img1, img2):
+    (_, channel, _, _) = img1.size()
+    window_size = 11
+    window = create_window(window_size, channel)
+
+    if img1.is_cuda:
+        window = window.cuda(img1.get_device())
+    window = window.type_as(img1)
+
+    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
+
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
+    return ssim_map.mean()
+
+
+def PSNR(img1, img2):
+    mse = np.mean((img1 / 255. - img2 / 255.) ** 2)
+    if mse == 0:
+        return 100
+    PIXEL_MAX = 1
+    return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
@@ -0,0 +1,48 @@
+from __future__ import print_function
+
+import numpy as np
+from PIL import Image
+import numpy as np
+import os
+import matplotlib.pyplot as plt
+import torch
+
+def load_image(path):
+    if(path[-3:] == 'dng'):
+        import rawpy
+        with rawpy.imread(path) as raw:
+            img = raw.postprocess()
+    elif(path[-3:]=='bmp' or path[-3:]=='jpg' or path[-3:]=='png'):
+        import cv2
+        return cv2.imread(path)[:,:,::-1]
+    else:
+        img = (255*plt.imread(path)[:,:,:3]).astype('uint8')
+
+    return img
+
+def save_image(image_numpy, image_path, ):
+    image_pil = Image.fromarray(image_numpy)
+    image_pil.save(image_path)
+
+def mkdirs(paths):
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
+# def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=1.):
+    image_numpy = image_tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
+    return image_numpy.astype(imtype)
+
+def im2tensor(image, imtype=np.uint8, cent=1., factor=255./2.):
+# def im2tensor(image, imtype=np.uint8, cent=1., factor=1.):
+    return torch.Tensor((image / factor - cent)
+                        [:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
@@ -0,0 +1,216 @@
+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)