text_embeddings.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import torch
import os
import numpy as np
import clip
import argparse


## VOC12 w/ background
VOC12_wbg_classes = ['background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
               'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
               'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa',
               'train', 'tvmonitor']


VOC12_wbg_classes_w_concepts4 = [
        ['background', 'bed', 'building', 'cabinet', 'ceiling', 'curtain', 'door', 'fence',
         'floor', 'grass', 'ground', 'mountain', 'road', 'rock', 'shelves', 'sidewalk', 'sky',
         'snow', 'tree', 'wall', 'water', 'window', 'hang glider', 'helicopter', 'jet ski',
         'go-cart', 'tractor', 'emergency vehicle', 'lorry', 'truck', 'lion', 'stool', 'bench',
         'wheelchair', 'coffee table', 'desk', 'side table', 'picnic bench', 'wolve',
         'flowers in a vase', 'goat', 'tram', 'laptop', 'advertising display', 'vehicle interior'],
        ['aeroplane', 'airplane', 'glider'],
        ['bicycle', 'tricycle', 'unicycle'],
        ['bird'],
        ['boat', 'ship', 'rowing boat', 'pedalo'],
        ['bottle', 'plastic bottle', 'glass bottle', 'feeding bottle'],
        ['bus', 'minibus'],
        ['car', 'van', 'large family car', 'realistic toy car'],
        ['cat', 'domestic cat'],
        ['chair', 'armchair', 'deckchair'],
        ['cow'],
        ['dining table', 'table for eating at'],
        ['dog', 'domestic dog'],
        ['horse', 'pony', 'donkey', 'mule'],
        ['motorbike', 'moped', 'scooter', 'sidecar'],
        ['person', 'people', 'baby', 'face'],
        ['potted plant', 'indoor plant in a pot', 'outdoor plant in a pot'],
        ['sheep'],
        ['sofa'],
        ['train', 'train carriage'],
        ['tv', 'monitor', 'standalone screen'],
]


## COCO Stuff
COCO_stuff_classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
        'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
        'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
        'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
        'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
        'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
        'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
        'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
        'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
        'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
        'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
        'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
        'scissors', 'teddy bear', 'hair drier', 'toothbrush', 'banner',
        'blanket', 'branch', 'bridge', 'building-other', 'bush', 'cabinet',
        'cage', 'cardboard', 'carpet', 'ceiling-other', 'ceiling-tile',
        'cloth', 'clothes', 'clouds', 'counter', 'cupboard', 'curtain',
        'desk-stuff', 'dirt', 'door-stuff', 'fence', 'floor-marble',
        'floor-other', 'floor-stone', 'floor-tile', 'floor-wood',
        'flower', 'fog', 'food-other', 'fruit', 'furniture-other', 'grass',
        'gravel', 'ground-other', 'hill', 'house', 'leaves', 'light', 'mat',
        'metal', 'mirror-stuff', 'moss', 'mountain', 'mud', 'napkin', 'net',
        'paper', 'pavement', 'pillow', 'plant-other', 'plastic', 'platform',
        'playingfield', 'railing', 'railroad', 'river', 'road', 'rock', 'roof',
        'rug', 'salad', 'sand', 'sea', 'shelf', 'sky-other', 'skyscraper',
        'snow', 'solid-other', 'stairs', 'stone', 'straw', 'structural-other',
        'table', 'tent', 'textile-other', 'towel', 'tree', 'vegetable',
        'wall-brick', 'wall-concrete', 'wall-other', 'wall-panel',
        'wall-stone', 'wall-tile', 'wall-wood', 'water-other', 'waterdrops',
        'window-blind', 'window-other', 'wood']

# COCO classes
COCO_classes = ['background', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
        'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
        'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
        'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
        'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
        'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
        'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
        'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
        'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
        'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
        'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
        'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
        'scissors', 'teddy bear', 'hair drier', 'toothbrush']

# ADE20k classes
ADE_classes = [
    "wall", "building", "sky", "floor", "tree", "ceiling", "road", "bed ", "windowpane", "grass", "cabinet",
    "sidewalk", "person", "earth", "door", "table", "mountain", "plant", "curtain", "chair", "car", "water",
    "painting", "sofa", "shelf", "house", "sea", "mirror", "rug", "field", "armchair", "seat", "fence", "desk",
    "rock", "wardrobe", "lamp", "bathtub", "railing", "cushion", "base", "box", "column", "signboard",
    "chest of drawers","counter", "sand", "sink", "skyscraper", "fireplace", "refrigerator", "grandstand",
    "path", "stairs", "runway", "case", "pool table", "pillow", "screen door", "stairway", "river", "bridge",
    "bookcase", "blind", "coffee table", "toilet", "flower", "book", "hill", "bench", "countertop", "stove",
    "palm", "kitchen island", "computer", "swivel chair", "boat", "bar", "arcade machine", "hovel", "bus",
    "towel", "light", "truck", "tower", "chandelier", "awning", "streetlight", "booth", "television receiver",
    "airplane", "dirt track", "apparel", "pole", "land", "bannister", "escalator", "ottoman", "bottle",
    "buffet", "poster", "stage", "van", "ship", "fountain", "conveyer belt", "canopy", "washer", "plaything",
    "swimming pool", "stool", "barrel", "basket", "waterfall", "tent", "bag", "minibike", "cradle", "oven",
    "ball", "food", "step", "tank", "trade name", "microwave", "pot", "animal", "bicycle", "lake", "dishwasher",
    "screen", "blanket", "sculpture", "hood", "sconce", "vase", "traffic light", "tray", "ashcan", "fan", "pier",
    "crt screen", "plate", "monitor", "bulletin board", "shower", "radiator", "glass", "clock", "flag"]

## Cityscapes
Cityscapes_classes = ['road', 'sidewalk', 'building', 'wall', 'fence', 'pole',
        'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky',
        'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle',
        'bicycle']



Cityscapes_classes_w_concepts3 = [
    ['road', 'street', 'parking space'],
    ['sidewalk'],
    ['building', 'skyscaper', 'house', 'bus stop building', 'garage', 'car port', 'scaffolding'],
    ['individual standing wall, which is not part of a building'],
    ['fence', 'hole in fence'],
    ['pole', 'sign pole', 'traffic light pole'],
    ['traffic light'],
    ['traffic sign', 'parking sign', 'direction sign'],
    ['vegetation', 'tree', 'hedge'],
    ['terrain', 'grass', 'soil', 'sand'],
    ['sky'],
    ['person', 'pedestrian', 'walking person', 'standing person', 'person sitting on the ground',
     'person sitting on a bench', 'person sitting on a chair'],
    ['rider', 'cyclist', 'motorcyclist'],
    ['car', 'jeep', 'SUV', 'van'],
    ['truck', 'box truck', 'pickup truck', 'truck trailer'],
    ['bus'],
    ['train', 'tram'],
    ['motorcycle', 'moped', 'scooter'],
    ['bicycle'],
]


def single_template(save_path, class_names, model):
    with torch.no_grad():
        texts = torch.cat([clip.tokenize(f"a photo of a {c}") for c in class_names]).cuda()
        text_embeddings = model.encode_text(texts)
        text_embeddings = text_embeddings / text_embeddings.norm(dim=-1, keepdim=True)
        np.save(save_path, text_embeddings.cpu().numpy())
    return text_embeddings

def single_template_concept_avg(save_path, class_concept_list, model):
    all_concepts = [j for sub in class_concept_list for j in sub]
    print(all_concepts)
    with torch.no_grad():
        tokens = torch.cat([clip.tokenize(f"a photo of a {c}") for c in all_concepts]).cuda()
        concept_embeddings = model.encode_text(tokens)
        print(concept_embeddings.shape)
        avg_concept_embeddings = []
        concept_idx = 0
        for i, concepts in enumerate(class_concept_list):
            n_concepts_for_class = len(concepts)
            print(concept_idx, concept_idx + n_concepts_for_class, all_concepts[concept_idx : concept_idx+n_concepts_for_class])
            concepts_for_class = concept_embeddings[concept_idx : concept_idx+n_concepts_for_class]
            assert concepts_for_class.shape[0] == n_concepts_for_class
            avg_concept = torch.sum(concepts_for_class, dim=0) / n_concepts_for_class
            avg_concept_embeddings.append(avg_concept)
            concept_idx += n_concepts_for_class
        avg_concept_embeddings = torch.stack(avg_concept_embeddings, dim=0)
        assert avg_concept_embeddings.shape[0] == len(class_concept_list)
        avg_concept_embeddings = avg_concept_embeddings / avg_concept_embeddings.norm(dim=-1, keepdim=True)
        if save_path is not None:
            np.save(save_path, avg_concept_embeddings.cpu().numpy())
    return avg_concept_embeddings

def aggregate_concept_predictions(pred, class_to_concept_idxs):
    B, _, H, W = pred.shape
    agg_pred = torch.zeros(B, len(class_to_concept_idxs), H, W, device=pred.device)
    for cls_i, conc_i in class_to_concept_idxs.items():
        agg_pred[:, cls_i] = pred[:, conc_i].max(dim=1).values
    return agg_pred

def flatten_class_concepts(class_concepts):
    concepts = []
    concept_to_class_idx = {}
    class_to_concept_idxs = {}
    for i, cls_concepts in enumerate(class_concepts):
        for concept in cls_concepts:
            concept_to_class_idx[len(concepts)] = i
            if i not in class_to_concept_idxs:
                class_to_concept_idxs[i] = []
            class_to_concept_idxs[i].append(len(concepts))
            concepts.append(concept)
    return concepts, concept_to_class_idx, class_to_concept_idxs

def get_class_to_concept_idxs(save_path):
    if save_path == 'configs/_base_/datasets/text_embedding/voc12_wbg_concept4_single.npy':
        _, _, class_to_concept_idxs = flatten_class_concepts(VOC12_wbg_classes_w_concepts4)
    elif save_path == 'configs/_base_/datasets/text_embedding/cityscapes_concept3_single.npy':
        _, _, class_to_concept_idxs = flatten_class_concepts(Cityscapes_classes_w_concepts3)
    else:
        raise ValueError(save_path)
    return class_to_concept_idxs

def prepare_text_embedding(save_path):
    assert os.path.isdir(os.path.dirname(save_path))
    # assert not os.path.isfile(save_path)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, _ = clip.load('ViT-B/16', device)

    ## VOC12 with background:
    if save_path == 'configs/_base_/datasets/text_embedding/voc12_wbg_single.npy':
        single_template(save_path, VOC12_wbg_classes, model)
    elif save_path == 'configs/_base_/datasets/text_embedding/voc12_wbg_conceptavg4_single.npy':
        single_template_concept_avg(save_path, VOC12_wbg_classes_w_concepts4, model)
    elif save_path == 'configs/_base_/datasets/text_embedding/voc12_wbg_concept4_single.npy':
        flat_concepts, _, _ = flatten_class_concepts(VOC12_wbg_classes_w_concepts4)
        single_template(save_path, flat_concepts, model)
    ## COCO:
    elif save_path == 'configs/_base_/datasets/text_embedding/coco_single.npy':
        single_template(save_path, COCO_classes, model)
    ## Cityscapes:
    elif save_path == 'configs/_base_/datasets/text_embedding/cityscapes_single.npy':
        single_template(save_path, Cityscapes_classes, model)
    elif save_path == 'configs/_base_/datasets/text_embedding/cityscapes_concept3_single.npy':
        flat_concepts, _, _ = flatten_class_concepts(Cityscapes_classes_w_concepts3)
        single_template(save_path, flat_concepts, model)
    elif save_path == 'configs/_base_/datasets/text_embedding/cityscapes_conceptavg3_single.npy':
        single_template_concept_avg(save_path, Cityscapes_classes_w_concepts3, model)
    ## ADE20k:
    elif save_path == 'configs/_base_/datasets/text_embedding/ade_single.npy':
        single_template(save_path, ADE_classes, model)
    else:
        raise NotImplementedError(save_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('name')
    args = parser.parse_args()
    prepare_text_embedding(f'configs/_base_/datasets/text_embedding/{args.name}.npy')