InvokeAI 2.2.4.5__py3-none-any.whl

vendored/clipseg/evaluation_utils.py

@@ -0,0 +1,292 @@
+from torch.functional import Tensor
+from general_utils import load_model
+from torch.utils.data import DataLoader
+import torch
+import numpy as np
+
+def denorm(img):
+
+    np_input = False
+    if isinstance(img, np.ndarray):
+        img = torch.from_numpy(img)
+        np_input = True
+
+    mean = torch.Tensor([0.485, 0.456, 0.406])
+    std = torch.Tensor([0.229, 0.224, 0.225])
+
+    img_denorm = (img*std[:,None,None]) + mean[:,None,None]
+
+    if np_input:
+        img_denorm = np.clip(img_denorm.numpy(), 0, 1)
+    else:
+        img_denorm = torch.clamp(img_denorm, 0, 1)
+
+    return img_denorm
+
+
+def norm(img):
+    mean = torch.Tensor([0.485, 0.456, 0.406])
+    std = torch.Tensor([0.229, 0.224, 0.225])
+    return (img - mean[:,None,None]) / std[:,None,None]
+
+
+def fast_iou_curve(p, g):
+    
+    g = g[p.sort().indices]
+    p = torch.sigmoid(p.sort().values)
+    
+    scores = []
+    vals = np.linspace(0, 1, 50)
+
+    for q in vals:
+
+        n = int(len(g) * q)
+
+        valid = torch.where(p > q)[0]
+        if len(valid) > 0:
+            n = int(valid[0])
+        else:
+            n = len(g)
+
+        fn = g[:n].sum()
+        tn = n - fn
+        tp = g[n:].sum()
+        fp = len(g) - n - tp
+
+        iou = tp / (tp + fn + fp)
+
+        precision = tp / (tp + fp)
+        recall = tp / (tp + fn)
+
+        scores += [iou]
+        
+    return vals, scores
+
+
+def fast_rp_curve(p, g):
+    
+    g = g[p.sort().indices]
+    p = torch.sigmoid(p.sort().values)
+    
+    precisions, recalls = [], []
+    vals = np.linspace(p.min(), p.max(), 250)
+
+    for q in p[::100000]:
+
+        n = int(len(g) * q)
+
+        valid = torch.where(p > q)[0]
+        if len(valid) > 0:
+            n = int(valid[0])
+        else:
+            n = len(g)
+
+        fn = g[:n].sum()
+        tn = n - fn
+        tp = g[n:].sum()
+        fp = len(g) - n - tp
+
+        iou = tp / (tp + fn + fp)
+
+        precision = tp / (tp + fp)
+        recall = tp / (tp + fn)
+
+        precisions += [precision]
+        recalls += [recall]
+        
+    return recalls, precisions
+
+
+# Image processing
+
+def img_preprocess(batch, blur=0, grayscale=False, center_context=None, rect=False, rect_color=(255,0,0), rect_width=2, 
+                   brightness=1.0, bg_fac=1, colorize=False, outline=False, image_size=224):
+    import cv2
+
+    rw = rect_width
+
+    out = []
+    for img, mask in zip(batch[1], batch[2]):
+
+        img = img.cpu() if isinstance(img, torch.Tensor) else torch.from_numpy(img)
+        mask = mask.cpu() if isinstance(mask, torch.Tensor) else torch.from_numpy(mask)
+        
+        img *= brightness
+        img_bl = img
+        if blur > 0: # best 5
+            img_bl = torch.from_numpy(cv2.GaussianBlur(img.permute(1,2,0).numpy(), (15, 15), blur)).permute(2,0,1)
+        
+        if grayscale:
+            img_bl = img_bl[1][None]
+        
+        #img_inp = img_ratio*img*mask + (1-img_ratio)*img_bl
+        # img_inp = img_ratio*img*mask + (1-img_ratio)*img_bl * (1-mask)
+        img_inp = img*mask + (bg_fac) * img_bl * (1-mask)
+
+        if rect:
+            _, bbox = crop_mask(img, mask, context=0.1)
+            img_inp[:, bbox[2]: bbox[3], max(0, bbox[0]-rw):bbox[0]+rw] = torch.tensor(rect_color)[:,None,None]
+            img_inp[:, bbox[2]: bbox[3], max(0, bbox[1]-rw):bbox[1]+rw] = torch.tensor(rect_color)[:,None,None]
+            img_inp[:, max(0, bbox[2]-1): bbox[2]+rw, bbox[0]:bbox[1]] = torch.tensor(rect_color)[:,None,None]
+            img_inp[:, max(0, bbox[3]-1): bbox[3]+rw, bbox[0]:bbox[1]] = torch.tensor(rect_color)[:,None,None]
+
+
+        if center_context is not None:
+            img_inp = object_crop(img_inp, mask, context=center_context, image_size=image_size)
+
+        if colorize:
+            img_gray = denorm(img)
+            img_gray = cv2.cvtColor(img_gray.permute(1,2,0).numpy(), cv2.COLOR_RGB2GRAY)
+            img_gray = torch.stack([torch.from_numpy(img_gray)]*3)
+            img_inp = torch.tensor([1,0.2,0.2])[:,None,None] * img_gray * mask + bg_fac * img_gray * (1-mask)
+            img_inp = norm(img_inp)
+
+        if outline:
+            cont = cv2.findContours(mask.byte().numpy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            outline_img = np.zeros(mask.shape, dtype=np.uint8)
+            cv2.drawContours(outline_img, cont[0], -1, thickness=5, color=(255, 255, 255))
+            outline_img = torch.stack([torch.from_numpy(outline_img)]*3).float() / 255.
+            img_inp = torch.tensor([1,0,0])[:,None,None] *  outline_img + denorm(img_inp) * (1- outline_img)
+            img_inp = norm(img_inp)
+
+        out += [img_inp]
+
+    return torch.stack(out)
+
+
+def object_crop(img, mask, context=0.0, square=False, image_size=224):
+    img_crop, bbox = crop_mask(img, mask, context=context, square=square)
+    img_crop = pad_to_square(img_crop, channel_dim=0)
+    img_crop = torch.nn.functional.interpolate(img_crop.unsqueeze(0), (image_size, image_size)).squeeze(0)
+    return img_crop
+    
+
+def crop_mask(img, mask, context=0.0, square=False):
+    
+    assert img.shape[1:] == mask.shape
+    
+    bbox = [mask.max(0).values.argmax(), mask.size(0) - mask.max(0).values.flip(0).argmax()]
+    bbox += [mask.max(1).values.argmax(), mask.size(1) - mask.max(1).values.flip(0).argmax()]
+    bbox = [int(x) for x in bbox]
+    
+    width, height = (bbox[3] - bbox[2]), (bbox[1] - bbox[0])
+
+    # square mask
+    if square:
+        bbox[0] = int(max(0, bbox[0] - context * height))
+        bbox[1] = int(min(mask.size(0), bbox[1] + context * height))
+        bbox[2] = int(max(0, bbox[2] - context * width))
+        bbox[3] = int(min(mask.size(1), bbox[3] + context * width))
+
+        width, height = (bbox[3] - bbox[2]), (bbox[1] - bbox[0])
+        if height > width:
+            bbox[2] = int(max(0, (bbox[2] - 0.5*height)))
+            bbox[3] = bbox[2] + height
+        else:
+            bbox[0] = int(max(0, (bbox[0] - 0.5*width)))
+            bbox[1] = bbox[0] + width
+    else:
+        bbox[0] = int(max(0, bbox[0] - context * height))
+        bbox[1] = int(min(mask.size(0), bbox[1] + context * height))
+        bbox[2] = int(max(0, bbox[2] - context * width))
+        bbox[3] = int(min(mask.size(1), bbox[3] + context * width))
+
+    width, height = (bbox[3] - bbox[2]), (bbox[1] - bbox[0])
+    img_crop = img[:, bbox[2]: bbox[3], bbox[0]: bbox[1]]
+    return img_crop, bbox
+
+
+def pad_to_square(img, channel_dim=2, fill=0):
+    """
+
+
+    add padding such that a squared image is returned """
+    
+    from torchvision.transforms.functional import pad
+
+    if channel_dim == 2:
+        img = img.permute(2, 0, 1)
+    elif channel_dim == 0:
+        pass
+    else:
+        raise ValueError('invalid channel_dim')
+
+    h, w = img.shape[1:]
+    pady1 = pady2 = padx1 = padx2 = 0
+
+    if h > w:
+        padx1 = (h - w) // 2
+        padx2 = h - w - padx1
+    elif w > h:
+        pady1 = (w - h) // 2
+        pady2 = w - h - pady1
+
+    img_padded = pad(img, padding=(padx1, pady1, padx2, pady2), padding_mode='constant')
+
+    if channel_dim == 2:
+        img_padded = img_padded.permute(1, 2, 0)
+
+    return img_padded
+
+
+# qualitative
+
+def split_sentence(inp, limit=9):
+    t_new, current_len = [], 0
+    for k, t in enumerate(inp.split(' ')):
+        current_len += len(t) + 1
+        t_new += [t+' ']
+        # not last
+        if current_len > limit and k != len(inp.split(' ')) - 1:
+            current_len = 0
+            t_new += ['\n']
+
+    t_new = ''.join(t_new)
+    return t_new
+    
+
+from matplotlib import pyplot as plt
+
+
+def plot(imgs, *preds, labels=None, scale=1, cmap=plt.cm.magma, aps=None, gt_labels=None, vmax=None):
+    
+    row_off = 0 if labels is None else 1
+    _, ax = plt.subplots(len(imgs) + row_off, 1 + len(preds), figsize=(scale * float(1 + 2*len(preds)), scale * float(len(imgs)*2)))
+    [a.axis('off') for a in ax.flatten()]
+    
+    if labels is not None:
+        for j in range(len(labels)):
+            t_new = split_sentence(labels[j], limit=6)
+            ax[0, 1+ j].text(0.5, 0.1, t_new, ha='center', fontsize=3+ 10*scale)
+
+
+    for i in range(len(imgs)):
+        ax[i + row_off,0].imshow(imgs[i])
+        for j in range(len(preds)):
+            img = preds[j][i][0].detach().cpu().numpy()
+
+            if gt_labels is not None and labels[j] == gt_labels[i]:
+                print(j, labels[j], gt_labels[i])
+                edgecolor = 'red'
+                if aps is not None:
+                    ax[i + row_off, 1 + j].text(30, 70, f'AP: {aps[i]:.3f}', color='red', fontsize=8)
+            else:
+                edgecolor = 'k'
+
+            rect = plt.Rectangle([0,0], img.shape[0], img.shape[1], facecolor="none", 
+                                 edgecolor=edgecolor, linewidth=3)
+            ax[i + row_off,1 + j].add_patch(rect)
+
+            if vmax is None:
+                this_vmax = 1 
+            elif vmax == 'per_prompt':
+                this_vmax = max([preds[j][_i][0].max() for _i in range(len(imgs))])
+            elif vmax == 'per_image':
+                this_vmax = max([preds[_j][i][0].max() for _j in range(len(preds))])
+
+            ax[i + row_off,1 + j].imshow(img, vmin=0, vmax=this_vmax, cmap=cmap)
+
+    
+            # ax[i,1 + j].imshow(preds[j][i][0].detach().cpu().numpy(), vmin=preds[j].min(), vmax=preds[j].max())
+    plt.tight_layout()
+    plt.subplots_adjust(wspace=0.05, hspace=0.05)    

vendored/clipseg/general_utils.py

@@ -0,0 +1,272 @@
+import json
+import inspect
+import torch
+import os
+import sys
+import yaml
+from shutil import copy, copytree
+from os.path import join, dirname, realpath, expanduser, isfile, isdir, basename
+
+
+class Logger(object):
+
+    def __getattr__(self, k):
+        return print
+
+log = Logger()
+
+def training_config_from_cli_args():
+    experiment_name = sys.argv[1]
+    experiment_id = int(sys.argv[2])
+
+    yaml_config = yaml.load(open(f'experiments/{experiment_name}'), Loader=yaml.SafeLoader)
+
+    config = yaml_config['configuration']
+    config = {**config, **yaml_config['individual_configurations'][experiment_id]}
+    config = AttributeDict(config)
+    return config
+
+
+def score_config_from_cli_args():
+    experiment_name = sys.argv[1]
+    experiment_id = int(sys.argv[2])
+    
+
+    yaml_config = yaml.load(open(f'experiments/{experiment_name}'), Loader=yaml.SafeLoader)
+
+    config = yaml_config['test_configuration_common']
+
+    if type(yaml_config['test_configuration']) == list:
+        test_id = int(sys.argv[3])
+        config = {**config, **yaml_config['test_configuration'][test_id]}
+    else:
+        config = {**config, **yaml_config['test_configuration']}
+
+    if 'test_configuration' in yaml_config['individual_configurations'][experiment_id]:
+        config = {**config, **yaml_config['individual_configurations'][experiment_id]['test_configuration']}
+
+    train_checkpoint_id = yaml_config['individual_configurations'][experiment_id]['name']
+
+    config = AttributeDict(config)
+    return config, train_checkpoint_id
+
+
+def get_from_repository(local_name, repo_files, integrity_check=None, repo_dir='~/dataset_repository', 
+                        local_dir='~/datasets'):
+    """ copies files from repository to local folder.
+    
+    repo_files: list of filenames or list of tuples [filename, target path] 
+
+    e.g. get_from_repository('MyDataset', [['data/dataset1.tar', 'other/path/ds03.tar'])
+    will create a folder 'MyDataset' in local_dir, and extract the content of
+    '<repo_dir>/data/dataset1.tar' to <local_dir>/MyDataset/other/path.
+     """
+
+    local_dir = realpath(join(expanduser(local_dir), local_name))
+
+    dataset_exists = True
+
+    # check if folder is available
+    if not isdir(local_dir):
+        dataset_exists = False
+
+    if integrity_check is not None:
+        try:
+            integrity_ok = integrity_check(local_dir)
+        except BaseException:
+            integrity_ok = False
+
+        if integrity_ok:
+            log.hint('Passed custom integrity check')
+        else:
+            log.hint('Custom integrity check failed')
+
+        dataset_exists = dataset_exists and integrity_ok
+
+    if not dataset_exists:
+
+        repo_dir = realpath(expanduser(repo_dir))
+
+        for i, filename in enumerate(repo_files):
+
+            if type(filename) == str:
+                origin, target = filename, filename
+                archive_target = join(local_dir, basename(origin))
+                extract_target = join(local_dir)
+            else:
+                origin, target = filename
+                archive_target = join(local_dir, dirname(target), basename(origin))
+                extract_target = join(local_dir, dirname(target))
+            
+            archive_origin = join(repo_dir, origin)
+
+            log.hint(f'copy: {archive_origin} to {archive_target}')
+
+            # make sure the path exists
+            os.makedirs(dirname(archive_target), exist_ok=True)
+
+            if os.path.isfile(archive_target):
+                # only copy if size differs
+                if os.path.getsize(archive_target) != os.path.getsize(archive_origin):
+                    log.hint(f'file exists but filesize differs: target {os.path.getsize(archive_target)} vs. origin {os.path.getsize(archive_origin)}')
+                    copy(archive_origin, archive_target)
+            else:
+                copy(archive_origin, archive_target)
+
+            extract_archive(archive_target, extract_target, noarchive_ok=True)
+
+            # concurrent processes might have deleted the file
+            if os.path.isfile(archive_target):
+                os.remove(archive_target)
+
+
+def extract_archive(filename, target_folder=None, noarchive_ok=False):
+    from subprocess import run, PIPE
+
+    if filename.endswith('.tgz') or filename.endswith('.tar'):
+        command = f'tar -xf {filename}'
+        command += f' -C {target_folder}' if target_folder is not None else ''
+    elif filename.endswith('.tar.gz'):
+        command = f'tar -xzf {filename}'
+        command += f' -C {target_folder}' if target_folder is not None else ''
+    elif filename.endswith('zip'):
+        command = f'unzip {filename}'
+        command += f' -d {target_folder}' if target_folder is not None else ''
+    else:
+        if noarchive_ok:
+            return
+        else:
+            raise ValueError(f'unsuppored file ending of {filename}')
+
+    log.hint(command)
+    result = run(command.split(), stdout=PIPE, stderr=PIPE)
+    if result.returncode != 0:
+        print(result.stdout, result.stderr)
+
+
+class AttributeDict(dict):
+    """ 
+    An extended dictionary that allows access to elements as atttributes and counts 
+    these accesses. This way, we know if some attributes were never used. 
+    """
+
+    def __init__(self, *args, **kwargs):
+        from collections import Counter
+        super().__init__(*args, **kwargs)
+        self.__dict__['counter'] = Counter()
+
+    def __getitem__(self, k):
+        self.__dict__['counter'][k] += 1
+        return super().__getitem__(k)
+
+    def __getattr__(self, k):
+        self.__dict__['counter'][k] += 1
+        return super().get(k)
+
+    def __setattr__(self, k, v):
+        return super().__setitem__(k, v)
+
+    def __delattr__(self, k, v):
+        return super().__delitem__(k, v)    
+
+    def unused_keys(self, exceptions=()):
+        return [k for k in super().keys() if self.__dict__['counter'][k] == 0 and k not in exceptions]
+
+    def assume_no_unused_keys(self, exceptions=()):
+        if len(self.unused_keys(exceptions=exceptions)) > 0:
+            log.warning('Unused keys:', self.unused_keys(exceptions=exceptions))
+
+
+def get_attribute(name):
+    import importlib
+
+    if name is None:
+        raise ValueError('The provided attribute is None')
+    
+    name_split = name.split('.')
+    mod = importlib.import_module('.'.join(name_split[:-1]))
+    return getattr(mod, name_split[-1])
+
+
+
+def filter_args(input_args, default_args):
+
+    updated_args = {k: input_args[k] if k in input_args else v for k, v in default_args.items()}
+    used_args = {k: v for k, v in input_args.items() if k in default_args}
+    unused_args = {k: v for k, v in input_args.items() if k not in default_args}
+
+    return AttributeDict(updated_args), AttributeDict(used_args), AttributeDict(unused_args)
+
+
+def load_model(checkpoint_id, weights_file=None, strict=True, model_args='from_config', with_config=False):
+
+    config = json.load(open(join('logs', checkpoint_id, 'config.json')))
+
+    if model_args != 'from_config' and type(model_args) != dict:
+        raise ValueError('model_args must either be "from_config" or a dictionary of values')
+
+    model_cls = get_attribute(config['model'])
+
+    # load model
+    if model_args == 'from_config':
+        _, model_args, _ = filter_args(config, inspect.signature(model_cls).parameters)
+
+    model = model_cls(**model_args)
+
+    if weights_file is None:
+        weights_file = realpath(join('logs', checkpoint_id, 'weights.pth'))
+    else:
+        weights_file = realpath(join('logs', checkpoint_id, weights_file))
+
+    if isfile(weights_file):
+        weights = torch.load(weights_file)
+        for _, w in weights.items():
+            assert not torch.any(torch.isnan(w)), 'weights contain NaNs'
+        model.load_state_dict(weights, strict=strict)
+    else:
+        raise FileNotFoundError(f'model checkpoint {weights_file} was not found')
+
+    if with_config:
+        return model, config
+    
+    return model
+
+
+class TrainingLogger(object):
+
+    def __init__(self, model, log_dir, config=None, *args):
+        super().__init__()
+        self.model = model
+        self.base_path = join(f'logs/{log_dir}') if log_dir is not None else None
+
+        os.makedirs('logs/', exist_ok=True)
+        os.makedirs(self.base_path, exist_ok=True)
+
+        if config is not None:
+            json.dump(config, open(join(self.base_path, 'config.json'), 'w'))
+
+    def iter(self, i, **kwargs):
+        if i % 100 == 0 and 'loss' in kwargs:
+            loss = kwargs['loss']
+            print(f'iteration {i}: loss {loss:.4f}')
+
+    def save_weights(self, only_trainable=False, weight_file='weights.pth'):
+        if self.model is None:
+            raise AttributeError('You need to provide a model reference when initializing TrainingTracker to save weights.')
+
+        weights_path = join(self.base_path, weight_file)
+
+        weight_dict = self.model.state_dict()
+
+        if only_trainable:
+            weight_dict = {n: weight_dict[n] for n, p in self.model.named_parameters() if p.requires_grad}
+        
+        torch.save(weight_dict, weights_path)
+        log.info(f'Saved weights to {weights_path}')
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """ automatically stop processes if used in a context manager """
+        pass