lama

1
import os
2

3
import cv2
4
import numpy as np
5

6
from saicinpainting.training.visualizers.base import BaseVisualizer, visualize_mask_and_images_batch
7
from saicinpainting.utils import check_and_warn_input_range
8

9

10
class DirectoryVisualizer(BaseVisualizer):
11
    DEFAULT_KEY_ORDER = 'image predicted_image inpainted'.split(' ')
12

13
    def __init__(self, outdir, key_order=DEFAULT_KEY_ORDER, max_items_in_batch=10,
14
                 last_without_mask=True, rescale_keys=None):
15
        self.outdir = outdir
16
        os.makedirs(self.outdir, exist_ok=True)
17
        self.key_order = key_order
18
        self.max_items_in_batch = max_items_in_batch
19
        self.last_without_mask = last_without_mask
20
        self.rescale_keys = rescale_keys
21

22
    def __call__(self, epoch_i, batch_i, batch, suffix='', rank=None):
23
        check_and_warn_input_range(batch['image'], 0, 1, 'DirectoryVisualizer target image')
24
        vis_img = visualize_mask_and_images_batch(batch, self.key_order, max_items=self.max_items_in_batch,
25
                                                  last_without_mask=self.last_without_mask,
26
                                                  rescale_keys=self.rescale_keys)
27

28
        vis_img = np.clip(vis_img * 255, 0, 255).astype('uint8')
29

30
        curoutdir = os.path.join(self.outdir, f'epoch{epoch_i:04d}{suffix}')
31
        os.makedirs(curoutdir, exist_ok=True)
32
        rank_suffix = f'_r{rank}' if rank is not None else ''
33
        out_fname = os.path.join(curoutdir, f'batch{batch_i:07d}{rank_suffix}.jpg')
34

35
        vis_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
36
        cv2.imwrite(out_fname, vis_img)
37