hierarcy_clip.py

# -*- coding: utf-8 -*-
"""Hierarcy_Clip.ipynb

Automatically generated by Colaboratory.

Original file is located at the same repo
    Hierarcy_Clip.ipynb

# CVPR2023 Hierarchy-CLIP: Improving Zero-shot Generalization and Robustness of Multi-modal Models

*Licensed under the Apache License, Version 2.0.*

This is a colab for Hierarchy-CLIP

## 1. Reproducing CLIP zero-shot ImageNet classification performance

Replicating the latest OpenAI [Colab](https://github.com/openai/CLIP/blob/main/notebooks/Prompt_Engineering_for_ImageNet.ipynb).
"""

# Commented out IPython magic to ensure Python compatibility.
# %pylab inline
import tensorflow_datasets as tfds
import jax
import jax.numpy as jnp
from tqdm import tqdm
import tensorflow as tf
import random
import os
import json
from scipy.special import softmax
from PIL import Image
import pandas as pd
from scenic.projects.baselines.clip import model as clip
from scenic.projects.baselines.clip import tokenizer as clip_tokenizer

#@title ImageNet classNames
# https://github.com/openai/CLIP/blob/main/notebooks/Prompt_Engineering_for_ImageNet.ipynb
clip.IMAGENET_CLASSES = ['tench', 'goldfish', 'great white shark', 'tiger shark', 'hammerhead shark', 'electric ray', 'stingray', 'rooster', 'hen', 'ostrich', 'brambling', 'goldfinch', 'house finch', 'junco', 'indigo bunting', 'American robin', 'bulbul', 'jay', 'magpie', 'chickadee', 'American dipper', 'kite (bird of prey)', 'bald eagle', 'vulture', 'great grey owl', 'fire salamander', 'smooth newt', 'newt', 'spotted salamander', 'axolotl', 'American bullfrog', 'tree frog', 'tailed frog', 'loggerhead sea turtle', 'leatherback sea turtle', 'mud turtle', 'terrapin', 'box turtle', 'banded gecko', 'green iguana', 'Carolina anole', 'desert grassland whiptail lizard', 'agama', 'frilled-necked lizard', 'alligator lizard', 'Gila monster', 'European green lizard', 'chameleon', 'Komodo dragon', 'Nile crocodile', 'American alligator', 'triceratops', 'worm snake', 'ring-necked snake', 'eastern hog-nosed snake', 'smooth green snake', 'kingsnake', 'garter snake', 'water snake', 'vine snake', 'night snake', 'boa constrictor', 'African rock python', 'Indian cobra', 'green mamba', 'sea snake', 'Saharan horned viper', 'eastern diamondback rattlesnake', 'sidewinder rattlesnake', 'trilobite', 'harvestman', 'scorpion', 'yellow garden spider', 'barn spider', 'European garden spider', 'southern black widow', 'tarantula', 'wolf spider', 'tick', 'centipede', 'black grouse', 'ptarmigan', 'ruffed grouse', 'prairie grouse', 'peafowl', 'quail', 'partridge', 'african grey parrot', 'macaw', 'sulphur-crested cockatoo', 'lorikeet', 'coucal', 'bee eater', 'hornbill', 'hummingbird', 'jacamar', 'toucan', 'duck', 'red-breasted merganser', 'goose', 'black swan', 'tusker', 'echidna', 'platypus', 'wallaby', 'koala', 'wombat', 'jellyfish', 'sea anemone', 'brain coral', 'flatworm', 'nematode', 'conch', 'snail', 'slug', 'sea slug', 'chiton', 'chambered nautilus', 'Dungeness crab', 'rock crab', 'fiddler crab', 'red king crab', 'American lobster', 'spiny lobster', 'crayfish', 'hermit crab', 'isopod', 'white stork', 'black stork', 'spoonbill', 'flamingo', 'little blue heron', 'great egret', 'bittern bird', 'crane bird', 'limpkin', 'common gallinule', 'American coot', 'bustard', 'ruddy turnstone', 'dunlin', 'common redshank', 'dowitcher', 'oystercatcher', 'pelican', 'king penguin', 'albatross', 'grey whale', 'killer whale', 'dugong', 'sea lion', 'Chihuahua', 'Japanese Chin', 'Maltese', 'Pekingese', 'Shih Tzu', 'King Charles Spaniel', 'Papillon', 'toy terrier', 'Rhodesian Ridgeback', 'Afghan Hound', 'Basset Hound', 'Beagle', 'Bloodhound', 'Bluetick Coonhound', 'Black and Tan Coonhound', 'Treeing Walker Coonhound', 'English foxhound', 'Redbone Coonhound', 'borzoi', 'Irish Wolfhound', 'Italian Greyhound', 'Whippet', 'Ibizan Hound', 'Norwegian Elkhound', 'Otterhound', 'Saluki', 'Scottish Deerhound', 'Weimaraner', 'Staffordshire Bull Terrier', 'American Staffordshire Terrier', 'Bedlington Terrier', 'Border Terrier', 'Kerry Blue Terrier', 'Irish Terrier', 'Norfolk Terrier', 'Norwich Terrier', 'Yorkshire Terrier', 'Wire Fox Terrier', 'Lakeland Terrier', 'Sealyham Terrier', 'Airedale Terrier', 'Cairn Terrier', 'Australian Terrier', 'Dandie Dinmont Terrier', 'Boston Terrier', 'Miniature Schnauzer', 'Giant Schnauzer', 'Standard Schnauzer', 'Scottish Terrier', 'Tibetan Terrier', 'Australian Silky Terrier', 'Soft-coated Wheaten Terrier', 'West Highland White Terrier', 'Lhasa Apso', 'Flat-Coated Retriever', 'Curly-coated Retriever', 'Golden Retriever', 'Labrador Retriever', 'Chesapeake Bay Retriever', 'German Shorthaired Pointer', 'Vizsla', 'English Setter', 'Irish Setter', 'Gordon Setter', 'Brittany dog', 'Clumber Spaniel', 'English Springer Spaniel', 'Welsh Springer Spaniel', 'Cocker Spaniel', 'Sussex Spaniel', 'Irish Water Spaniel', 'Kuvasz', 'Schipperke', 'Groenendael dog', 'Malinois', 'Briard', 'Australian Kelpie', 'Komondor', 'Old English Sheepdog', 'Shetland Sheepdog', 'collie', 'Border Collie', 'Bouvier des Flandres dog', 'Rottweiler', 'German Shepherd Dog', 'Dobermann', 'Miniature Pinscher', 'Greater Swiss Mountain Dog', 'Bernese Mountain Dog', 'Appenzeller Sennenhund', 'Entlebucher Sennenhund', 'Boxer', 'Bullmastiff', 'Tibetan Mastiff', 'French Bulldog', 'Great Dane', 'St. Bernard', 'husky', 'Alaskan Malamute', 'Siberian Husky', 'Dalmatian', 'Affenpinscher', 'Basenji', 'pug', 'Leonberger', 'Newfoundland dog', 'Great Pyrenees dog', 'Samoyed', 'Pomeranian', 'Chow Chow', 'Keeshond', 'brussels griffon', 'Pembroke Welsh Corgi', 'Cardigan Welsh Corgi', 'Toy Poodle', 'Miniature Poodle', 'Standard Poodle', 'Mexican hairless dog (xoloitzcuintli)', 'grey wolf', 'Alaskan tundra wolf', 'red wolf or maned wolf', 'coyote', 'dingo', 'dhole', 'African wild dog', 'hyena', 'red fox', 'kit fox', 'Arctic fox', 'grey fox', 'tabby cat', 'tiger cat', 'Persian cat', 'Siamese cat', 'Egyptian Mau', 'cougar', 'lynx', 'leopard', 'snow leopard', 'jaguar', 'lion', 'tiger', 'cheetah', 'brown bear', 'American black bear', 'polar bear', 'sloth bear', 'mongoose', 'meerkat', 'tiger beetle', 'ladybug', 'ground beetle', 'longhorn beetle', 'leaf beetle', 'dung beetle', 'rhinoceros beetle', 'weevil', 'fly', 'bee', 'ant', 'grasshopper', 'cricket insect', 'stick insect', 'cockroach', 'praying mantis', 'cicada', 'leafhopper', 'lacewing', 'dragonfly', 'damselfly', 'red admiral butterfly', 'ringlet butterfly', 'monarch butterfly', 'small white butterfly', 'sulphur butterfly', 'gossamer-winged butterfly', 'starfish', 'sea urchin', 'sea cucumber', 'cottontail rabbit', 'hare', 'Angora rabbit', 'hamster', 'porcupine', 'fox squirrel', 'marmot', 'beaver', 'guinea pig', 'common sorrel horse', 'zebra', 'pig', 'wild boar', 'warthog', 'hippopotamus', 'ox', 'water buffalo', 'bison', 'ram (adult male sheep)', 'bighorn sheep', 'Alpine ibex', 'hartebeest', 'impala (antelope)', 'gazelle', 'arabian camel', 'llama', 'weasel', 'mink', 'European polecat', 'black-footed ferret', 'otter', 'skunk', 'badger', 'armadillo', 'three-toed sloth', 'orangutan', 'gorilla', 'chimpanzee', 'gibbon', 'siamang', 'guenon', 'patas monkey', 'baboon', 'macaque', 'langur', 'black-and-white colobus', 'proboscis monkey', 'marmoset', 'white-headed capuchin', 'howler monkey', 'titi monkey', "Geoffroy's spider monkey", 'common squirrel monkey', 'ring-tailed lemur', 'indri', 'Asian elephant', 'African bush elephant', 'red panda', 'giant panda', 'snoek fish', 'eel', 'silver salmon', 'rock beauty fish', 'clownfish', 'sturgeon', 'gar fish', 'lionfish', 'pufferfish', 'abacus', 'abaya', 'academic gown', 'accordion', 'acoustic guitar', 'aircraft carrier', 'airliner', 'airship', 'altar', 'ambulance', 'amphibious vehicle', 'analog clock', 'apiary', 'apron', 'trash can', 'assault rifle', 'backpack', 'bakery', 'balance beam', 'balloon', 'ballpoint pen', 'Band-Aid', 'banjo', 'baluster / handrail', 'barbell', 'barber chair', 'barbershop', 'barn', 'barometer', 'barrel', 'wheelbarrow', 'baseball', 'basketball', 'bassinet', 'bassoon', 'swimming cap', 'bath towel', 'bathtub', 'station wagon', 'lighthouse', 'beaker', 'military hat (bearskin or shako)', 'beer bottle', 'beer glass', 'bell tower', 'baby bib', 'tandem bicycle', 'bikini', 'ring binder', 'binoculars', 'birdhouse', 'boathouse', 'bobsleigh', 'bolo tie', 'poke bonnet', 'bookcase', 'bookstore', 'bottle cap', 'hunting bow', 'bow tie', 'brass memorial plaque', 'bra', 'breakwater', 'breastplate', 'broom', 'bucket', 'buckle', 'bulletproof vest', 'high-speed train', 'butcher shop', 'taxicab', 'cauldron', 'candle', 'cannon', 'canoe', 'can opener', 'cardigan', 'car mirror', 'carousel', 'tool kit', 'cardboard box / carton', 'car wheel', 'automated teller machine', 'cassette', 'cassette player', 'castle', 'catamaran', 'CD player', 'cello', 'mobile phone', 'chain', 'chain-link fence', 'chain mail', 'chainsaw', 'storage chest', 'chiffonier', 'bell or wind chime', 'china cabinet', 'Christmas stocking', 'church', 'movie theater', 'cleaver', 'cliff dwelling', 'cloak', 'clogs', 'cocktail shaker', 'coffee mug', 'coffeemaker', 'spiral or coil', 'combination lock', 'computer keyboard', 'candy store', 'container ship', 'convertible', 'corkscrew', 'cornet', 'cowboy boot', 'cowboy hat', 'cradle', 'construction crane', 'crash helmet', 'crate', 'infant bed', 'Crock Pot', 'croquet ball', 'crutch', 'cuirass', 'dam', 'desk', 'desktop computer', 'rotary dial telephone', 'diaper', 'digital clock', 'digital watch', 'dining table', 'dishcloth', 'dishwasher', 'disc brake', 'dock', 'dog sled', 'dome', 'doormat', 'drilling rig', 'drum', 'drumstick', 'dumbbell', 'Dutch oven', 'electric fan', 'electric guitar', 'electric locomotive', 'entertainment center', 'envelope', 'espresso machine', 'face powder', 'feather boa', 'filing cabinet', 'fireboat', 'fire truck', 'fire screen', 'flagpole', 'flute', 'folding chair', 'football helmet', 'forklift', 'fountain', 'fountain pen', 'four-poster bed', 'freight car', 'French horn', 'frying pan', 'fur coat', 'garbage truck', 'gas mask or respirator', 'gas pump', 'goblet', 'go-kart', 'golf ball', 'golf cart', 'gondola', 'gong', 'gown', 'grand piano', 'greenhouse', 'radiator grille', 'grocery store', 'guillotine', 'hair clip', 'hair spray', 'half-track', 'hammer', 'hamper', 'hair dryer', 'hand-held computer', 'handkerchief', 'hard disk drive', 'harmonica', 'harp', 'combine harvester', 'hatchet', 'holster', 'home theater', 'honeycomb', 'hook', 'hoop skirt', 'gymnastic horizontal bar', 'horse-drawn vehicle', 'hourglass', 'iPod', 'clothes iron', 'carved pumpkin', 'jeans', 'jeep', 'T-shirt', 'jigsaw puzzle', 'rickshaw', 'joystick', 'kimono', 'knee pad', 'knot', 'lab coat', 'ladle', 'lampshade', 'laptop computer', 'lawn mower', 'lens cap', 'letter opener', 'library', 'lifeboat', 'lighter', 'limousine', 'ocean liner', 'lipstick', 'slip-on shoe', 'lotion', 'music speaker', 'loupe magnifying glass', 'sawmill', 'magnetic compass', 'messenger bag', 'mailbox', 'tights', 'one-piece bathing suit', 'manhole cover', 'maraca', 'marimba', 'mask', 'matchstick', 'maypole', 'maze', 'measuring cup', 'medicine cabinet', 'megalith', 'microphone', 'microwave oven', 'military uniform', 'milk can', 'minibus', 'miniskirt', 'minivan', 'missile', 'mitten', 'mixing bowl', 'mobile home', 'ford model t', 'modem', 'monastery', 'monitor', 'moped', 'mortar and pestle', 'graduation cap', 'mosque', 'mosquito net', 'vespa', 'mountain bike', 'tent', 'computer mouse', 'mousetrap', 'moving van', 'muzzle', 'metal nail', 'neck brace', 'necklace', 'baby pacifier', 'notebook computer', 'obelisk', 'oboe', 'ocarina', 'odometer', 'oil filter', 'pipe organ', 'oscilloscope', 'overskirt', 'bullock cart', 'oxygen mask', 'product packet / packaging', 'paddle', 'paddle wheel', 'padlock', 'paintbrush', 'pajamas', 'palace', 'pan flute', 'paper towel', 'parachute', 'parallel bars', 'park bench', 'parking meter', 'railroad car', 'patio', 'payphone', 'pedestal', 'pencil case', 'pencil sharpener', 'perfume', 'Petri dish', 'photocopier', 'plectrum', 'Pickelhaube', 'picket fence', 'pickup truck', 'pier', 'piggy bank', 'pill bottle', 'pillow', 'ping-pong ball', 'pinwheel', 'pirate ship', 'drink pitcher', 'block plane', 'planetarium', 'plastic bag', 'plate rack', 'farm plow', 'plunger', 'Polaroid camera', 'pole', 'police van', 'poncho', 'pool table', 'soda bottle', 'plant pot', "potter's wheel", 'power drill', 'prayer rug', 'printer', 'prison', 'missile', 'projector', 'hockey puck', 'punching bag', 'purse', 'quill', 'quilt', 'race car', 'racket', 'radiator', 'radio', 'radio telescope', 'rain barrel', 'recreational vehicle', 'fishing casting reel', 'reflex camera', 'refrigerator', 'remote control', 'restaurant', 'revolver', 'rifle', 'rocking chair', 'rotisserie', 'eraser', 'rugby ball', 'ruler measuring stick', 'sneaker', 'safe', 'safety pin', 'salt shaker', 'sandal', 'sarong', 'saxophone', 'scabbard', 'weighing scale', 'school bus', 'schooner', 'scoreboard', 'CRT monitor', 'screw', 'screwdriver', 'seat belt', 'sewing machine', 'shield', 'shoe store', 'shoji screen / room divider', 'shopping basket', 'shopping cart', 'shovel', 'shower cap', 'shower curtain', 'ski', 'balaclava ski mask', 'sleeping bag', 'slide rule', 'sliding door', 'slot machine', 'snorkel', 'snowmobile', 'snowplow', 'soap dispenser', 'soccer ball', 'sock', 'solar thermal collector', 'sombrero', 'soup bowl', 'keyboard space bar', 'space heater', 'space shuttle', 'spatula', 'motorboat', 'spider web', 'spindle', 'sports car', 'spotlight', 'stage', 'steam locomotive', 'through arch bridge', 'steel drum', 'stethoscope', 'scarf', 'stone wall', 'stopwatch', 'stove', 'strainer', 'tram', 'stretcher', 'couch', 'stupa', 'submarine', 'suit', 'sundial', 'sunglasses', 'sunglasses', 'sunscreen', 'suspension bridge', 'mop', 'sweatshirt', 'swim trunks / shorts', 'swing', 'electrical switch', 'syringe', 'table lamp', 'tank', 'tape player', 'teapot', 'teddy bear', 'television', 'tennis ball', 'thatched roof', 'front curtain', 'thimble', 'threshing machine', 'throne', 'tile roof', 'toaster', 'tobacco shop', 'toilet seat', 'torch', 'totem pole', 'tow truck', 'toy store', 'tractor', 'semi-trailer truck', 'tray', 'trench coat', 'tricycle', 'trimaran', 'tripod', 'triumphal arch', 'trolleybus', 'trombone', 'hot tub', 'turnstile', 'typewriter keyboard', 'umbrella', 'unicycle', 'upright piano', 'vacuum cleaner', 'vase', 'vaulted or arched ceiling', 'velvet fabric', 'vending machine', 'vestment', 'viaduct', 'violin', 'volleyball', 'waffle iron', 'wall clock', 'wallet', 'wardrobe', 'military aircraft', 'sink', 'washing machine', 'water bottle', 'water jug', 'water tower', 'whiskey jug', 'whistle', 'hair wig', 'window screen', 'window shade', 'Windsor tie', 'wine bottle', 'airplane wing', 'wok', 'wooden spoon', 'wool', 'split-rail fence', 'shipwreck', 'sailboat', 'yurt', 'website', 'comic book', 'crossword', 'traffic or street sign', 'traffic light', 'dust jacket', 'menu', 'plate', 'guacamole', 'consomme', 'hot pot', 'trifle', 'ice cream', 'popsicle', 'baguette', 'bagel', 'pretzel', 'cheeseburger', 'hot dog', 'mashed potatoes', 'cabbage', 'broccoli', 'cauliflower', 'zucchini', 'spaghetti squash', 'acorn squash', 'butternut squash', 'cucumber', 'artichoke', 'bell pepper', 'cardoon', 'mushroom', 'Granny Smith apple', 'strawberry', 'orange', 'lemon', 'fig', 'pineapple', 'banana', 'jackfruit', 'cherimoya (custard apple)', 'pomegranate', 'hay', 'carbonara', 'chocolate syrup', 'dough', 'meatloaf', 'pizza', 'pot pie', 'burrito', 'red wine', 'espresso', 'tea cup', 'eggnog', 'mountain', 'bubble', 'cliff', 'coral reef', 'geyser', 'lakeshore', 'promontory', 'sandbar', 'beach', 'valley', 'volcano', 'baseball player', 'bridegroom', 'scuba diver', 'rapeseed', 'daisy', "yellow lady's slipper", 'corn', 'acorn', 'rose hip', 'horse chestnut seed', 'coral fungus', 'agaric', 'gyromitra', 'stinkhorn mushroom', 'earth star fungus', 'hen of the woods mushroom', 'bolete', 'corn cob', 'toilet paper']

model_name = 'vit_b16' # we could change different backbone

model = clip.MODELS[model_name]()
vars = clip.load_model_vars(model_name)

encode_text = jax.jit(lambda texts: model.apply(vars, texts, method=model.encode_text))
encode_image = jax.jit(lambda x: model.apply(vars, x, method=model.encode_image))

tokenize_fn = clip_tokenizer.build_tokenizer()

def permute_words(text):
  words = text.split(' ')
  random.shuffle(words)
  return ' '.join(words)

def zeroshot_classifier(classnames, templates, permute=False):
  zeroshot_weights = []
  permute_fn = permute_words if permute else lambda x: x
  for classname in tqdm(classnames):
    texts = [permute_fn(template.format(classname)) for template in templates]
    class_embeddings = encode_text(tokenize_fn(texts))
    class_embedding = class_embeddings.mean(0)
    class_embedding /= jnp.linalg.norm(class_embedding)
    zeroshot_weights.append(class_embedding)
  return jnp.stack(zeroshot_weights, axis=1)

# Readout weights with prompt engineering
weights_prompteng = zeroshot_classifier(clip.IMAGENET_CLASSES, clip.PROMPTS)

# Readout weights with modified ImageNet class names only
weights_name = zeroshot_classifier(clip.IMAGENET_CLASSES, ['{}'])

def preprocess(batch, size=224):
  batch = tf.image.convert_image_dtype(batch, dtype=tf.float32)
  return central_crop(resize_small(batch, size), (size, size))

def central_crop(image, crop_size):
    '''
    image
    crop_size: (h, w)
    '''
    h, w = crop_size[0], crop_size[1]
    dy = (tf.shape(image)[0] - h) // 2
    dx = (tf.shape(image)[1] - w) // 2
    return tf.image.crop_to_bounding_box(image, dy, dx, h, w)

def resize_small(image, smaller_size, method="area", antialias=True):
    '''
    image
    smaller_size: an integer, that represents a new size of the smaller side of
      an input image.
    '''
    h, w = tf.shape(image)[0], tf.shape(image)[1]

    # Figure out the necessary h/w.
    ratio = (
        tf.cast(smaller_size, tf.float32) /
        tf.cast(tf.minimum(h, w), tf.float32))
    h = tf.cast(tf.round(tf.cast(h, tf.float32) * ratio), tf.int32)
    w = tf.cast(tf.round(tf.cast(w, tf.float32) * ratio), tf.int32)

    dtype = image.dtype
    image = tf.image.resize(image, (h, w), method, antialias)
    return tf.cast(image, dtype)


def normalize(img):
  return (img - clip.IMAGE_MEAN) / clip.IMAGE_STD

def unnormalize(x):
  return x * clip.IMAGE_STD + clip.IMAGE_MEAN

def load_dataset(dataset='imagenet2012', split='validation', batch_size=1024):
  ds = tfds.load(dataset, split=split)
  def _preprocess(d):
    d['image'] = normalize(preprocess(d['image']))
    return d
  def _prepare(d):
    return jax.tree_map(lambda x: x._numpy(), d)
  batched_dataset = ds.map(_preprocess).batch(batch_size)
  batched_dataset = map(_prepare, batched_dataset)
  return batched_dataset

def load_dataset_info(dataset='imagenet2012', split='validation', batch_size=1024):
  ds, info = tfds.load(dataset, split=split, with_info="true")
  def _preprocess(d):
    d['image'] = normalize(preprocess(d['image']))
    return d
  def _prepare(d):
    return jax.tree_map(lambda x: x._numpy(), d)
  batched_dataset = ds.map(_preprocess).batch(batch_size)
  batched_dataset = map(_prepare, batched_dataset)
  return batched_dataset, info

def load_dataset_from(data_dir='YOUR/LOCAL/PATH/imagenet2012', dataset='imagenet2012', split='validation', batch_size=1024):
  # ds = tfds.load(dataset, split=split, data_dir=data_dir)
  ds = tfds.builder_from_directory(data_dir)
  ds = ds.as_dataset(split='validation')
  def _preprocess(d):
    d['image'] = normalize(preprocess(d['image']))
    return d
  def _prepare(d):
    return jax.tree_map(lambda x: x._numpy(), d)
  batched_dataset = ds.map(_preprocess).batch(batch_size)
  batched_dataset = map(_prepare, batched_dataset)
  return batched_dataset

def compute_image_embeddings(dset, norm_image=True):
  embeddings = []
  labels = []
  for batch in tqdm(dset):
    image_embedding = encode_image(batch['image'])
    if norm_image:
      image_embedding /= jnp.linalg.norm(image_embedding)
    embeddings.append(image_embedding)
    labels.append(batch['label'])
  return jnp.vstack(embeddings), jnp.hstack(labels)

def compute_accuracy(logits, labels):
  top_probs, top_labels = jax.lax.top_k(logits, 5)
  top1 = 100 * jnp.mean(top_labels[:, 0] == labels)
  top5 = 100 * jnp.sum(top_labels == labels[:, None]) / labels.shape[0]
  return top1, top5

dset = load_dataset('imagenet2012')
# You could also first download ImageNet and then process them with tensorflow_datasets and load them with function:
# dset = load_dataset_from(data_dir='YOUR/LOCAL/PATH/imagenet2012', dataset='imagenet2012', split='validation')
embeddings, labels = compute_image_embeddings(dset)

# More accurate and consistent result
logits_prompteng = np.matmul(embeddings, weights_prompteng)
logits_name = np.matmul(embeddings, weights_name)
top1_prompt, top5_prompt = compute_accuracy(logits_prompteng, labels)
top1_name, top5_name = compute_accuracy(logits_name, labels)
print(f'Prompt Engineering: top1={top1_prompt:.2f}%, top5={top5_prompt:.2f}%')
print(f'Class Names: top1={top1_name:.2f}%, top5={top5_name:.2f}%')

"""## 2. Build hierarchy from WordNet

The wordnet hierarchy is based on Github repo: https://github.com/niharikajainn/imagenet-ancestors-descendants

### WordNet parse
"""

root_path = '/YOUR/PATH/OF/DOWNLOADED/WORDNET/FILES' # VM yunhao

words_map = {}
child_map = {}
parent_map = {}
gloss_map = {} # description
# BEGIN GOOGLE-INTERNAL
words_path = os.path.join(root_path, "imagenet-ancestors-descendants", "words.txt")
gloss_path = os.path.join(root_path, "imagenet-ancestors-descendants", "gloss.txt")
child_map_path = os.path.join(root_path, "imagenet-ancestors-descendants", "is_a.txt")
imagenet_label_to_wordnet_file = os.path.join(root_path, "imagenet-ancestors-descendants", "imagenet_label_to_wordnet_synset.txt")
# END GOOGLE-INTERNAL

blank = ' '
comma_blank = ', '

#obtain wordnet_id mappings for all words
with tf.io.gfile.GFile(words_path, mode='r') as f:
	for line in f:
		line_split = line.split() # use blank " " to split
		wnid = line_split[0] # e.g., 'n03200357'
		words = line_split[1:] # e.g., ['electric,', 'electric', 'automobile,', 'electric', 'car'],
		words_map[wnid] = words
f.close()

#obtain wordnet_id mappings for all word description
with tf.io.gfile.GFile(gloss_path, mode='r') as f:
	for line in f:
		line_split = line.split()
		wnid = line_split[0]
		gloss = blank.join(line_split[1:])
		gloss_map[wnid] = gloss
f.close()

#obtain wordnet_id mappings for all parents-children
with tf.io.gfile.GFile(child_map_path, mode='r') as f:
	for line in f:
		parent, child = line.split()
		parent_map[child] = parent
		if parent not in child_map:
			child_map[parent] = [child]
		else:
			child_map[parent].append(child)
f.close()

# find details given wordnet_id
category = 'n02690373' #'n02084071' is dog
descendants = []
ancestors = []
# find Descendants and Ancestors
print(words_map[category])
print(gloss_map[category]+"\n")
#list all children
print("Descendants:\n")
if category in child_map:
  search = [child for child in child_map[category]]
while search: # go over all children (BFS)
  node = search.pop()
  print("\t"+ blank.join(words_map[node])+"\n")
  descendants.append(blank.join(words_map[node])) # keep all descendant
  if node in child_map: #has children
    [search.append(child) for child in child_map[node]]

#list all parents
print("Ancestors:\n")
if category in parent_map:
  node = parent_map[category] # only one parent class
else:
  node = category
while node in parent_map: # one way go up
  print("\t"+ blank.join(words_map[node])+"\n")
  ancestors.append(blank.join(words_map[node])) # keep all ancestor
  node = parent_map[node]
print("finish")

# get imagenet_id: wordnet_id mapping. e.g., '0': 'n01440764'
index_wdid = {}
with tf.io.gfile.GFile(imagenet_label_to_wordnet_file, mode='r') as f:
  for line in f:
    if "{'id'" in line:
      index_wdid[line.split(": {'")[0].split("{")[-1].split(" ")[-1]] = 'n' + line.split("-n")[0].split("'")[-1]      # {0: {'id': '01440764-n',

"""build hierarchy"""

# find ancestor, descendants and children for each class, may takes minites
an_des_dict_json = {}
for index in range(len(clip.IMAGENET_CLASSES)): # for 1000 classes
  category = index_wdid[str(index)]
  an_des_dict_json[str(index)] = {}
  descendants = []
  ancestors = []
  children = []
  # find Descendants and Ancestors
  #list all children
  if category in child_map:
    search = [child for child in child_map[category]] # here is only the child
    children = [blank.join(words_map[ele]) for ele in search]
  while search: # go over all children BFS priority queue
    node = search.pop()
    descendants.append(blank.join(words_map[node])) # keep all descendant
    if node in child_map: #has children
      [search.append(child) for child in child_map[node]]

  #list all parents
  node = parent_map[category] if category in parent_map else category
  while node in parent_map: # one way go up
    ancestors.append(blank.join(words_map[node])) # keep all ancestor
    node = parent_map[node]

  # save
  an_des_dict_json[str(index)]["wdid"] = category
  an_des_dict_json[str(index)]["words_map"] = blank.join(words_map[category])
  an_des_dict_json[str(index)]["clip_words_map"] = clip.IMAGENET_CLASSES[index]
  an_des_dict_json[str(index)]["ancestors"] = ancestors
  an_des_dict_json[str(index)]["descendants"] = descendants
  an_des_dict_json[str(index)]["children"] = children

"""### Utility functions for finding descendants and ancestors"""

# TODO(yunhaoge) Add examples for using the functions

def bottom_up_hierarchy(max_depth_ancestor=0, max_depth_descendant=1, use_children=True):
  """Obtain the class list by considering hierarchy
    here consisder only descendent because max_depth_ancestor=0
  """
  # Consider ancedescendant into clip zero-shot
  class_an_des_mapping = [] # start index of each imagenet class, used for following aggregation
  imagenet_classes_an_des = []
  class_an_des_mapping_list = [] # index details of each imagenet class
  i = 0
  for imagenet_idx, classname in enumerate(clip.IMAGENET_CLASSES):
    class_an_des_mapping.append(i)
    imagenet_classes_an_des.append(classname) # add original class name
    class_an_des_list = [] # relevant id for specific class
    i += 1 # add index first
    class_an_des_list.append(i)
    node_ancestors = an_des_dict_json[str(imagenet_idx)]["ancestors"]
    if use_children: # consider only children
      node_descendants = an_des_dict_json[str(imagenet_idx)]["children"]
    else:
      node_descendants = an_des_dict_json[str(imagenet_idx)]["descendants"]

    for an_idx, ancestor in enumerate(node_ancestors): # select ancestors
      if an_idx < max_depth_ancestor:
        if comma_blank in ancestor: # contains synonym, more than one, keep only one
          imagenet_classes_an_des.append(ancestor.split(comma_blank)[0])
          i += 1
          class_an_des_list.append(i)
        else:
          imagenet_classes_an_des.append(ancestor)
          i += 1
          class_an_des_list.append(i)
    for des_idx, descentant in enumerate(node_descendants): # select descendants
      if des_idx < max_depth_descendant:
        if comma_blank in descentant: # contains synonym, more than one, keep only one
          imagenet_classes_an_des.append(descentant.split(comma_blank)[0])
          i += 1
          class_an_des_list.append(i)
        else:
          imagenet_classes_an_des.append(descentant)
          i += 1
          class_an_des_list.append(i)
    class_an_des_mapping_list.append(class_an_des_list)
  return imagenet_classes_an_des, class_an_des_mapping, class_an_des_mapping_list

def top_down_hierarchy(interest_case, clip_templete, use_LCA = True, use_ancestor = False):
  """Obtain the word list by adding LCA."""
  # interest_case
  logits_ori_all = []
  for idx, interest_case_ele in enumerate(tqdm(interest_case)):
    # original word embedding
    word_list_ori = np.array(clip.IMAGENET_CLASSES)[np.array(interest_case_ele)].tolist()
    if use_LCA: # add LCA/A
      if use_ancestor: # use Ancestor
        word_list =  [word + blank + A_all[idx] for word in word_list_ori]
      else: # use LCA
        word_list =  [word + blank + LCA_all[idx] for word in word_list_ori]
    else:
      word_list =  [word for word in word_list_ori]
  return word_list

# fine the LCA for each image top 5
def find_LCA(top5_ancestor): # general
  """Obtain the Lowest comman ancestor of top 5 classes.

  Args:
      top5_ancestor: List of the ancestors for each candidate class
  Returns:
      LCA: Str, lowest comman ancestor
  """
  LCA = 'physical entity'
  # while (1) still have ancestor
  while min([len(ele) for ele in top5_ancestor]) > 0 :
    # find the highest for each class in topk
    current_roots = [topk_ancestor.pop() for topk_ancestor in top5_ancestor]
    current_roots_freq = Counter(current_roots)
    current_roots_freq = sorted(current_roots_freq.items(), key=lambda x: x[1], reverse=True) # become a list e.g.,[('ee', 2), ('ww', 1), ('cc', 1)]
    majority, majority_freq = current_roots_freq[0]
    if majority_freq == 5:
      LCA = majority if comma_blank not in majority else majority.split(comma_blank)[0]
  return LCA

"""## Uncertainty Estimation"""

# with children
# hierarchy ImageNet class work on same h for each class
max_depth_ancestor = 0 # here only consider descendent
max_depth_children = 10 # could change
imagenet_bottom_up_first_index_list = []
imagenet_bottom_up_all_classes = []
imagenet_bottom_up_mapping_index_list = [] # add list for analysis
i = 0
for imagenet_idx, classname in enumerate(clip.IMAGENET_CLASSES):
  imagenet_bottom_up_first_index_list.append(i) # index of each imagenet class
  imagenet_bottom_up_all_classes.append(classname) # add original class name
  class_an_des_list = [] # relevant id for specific class
  i += 1 # add index first
  class_an_des_list.append(i)
  node_ancestors = an_des_dict_json[str(imagenet_idx)]["ancestors"]
  node_children = an_des_dict_json[str(imagenet_idx)]["children"]

  for an_idx, ancestor in enumerate(node_ancestors): # select ancestors
    if an_idx < max_depth_ancestor:
      if ", " in ancestor: # contains synonym, more than one, keep only one
        imagenet_bottom_up_all_classes.append(ancestor.split(", ")[0])
        i += 1
        class_an_des_list.append(i)
      else:
        imagenet_bottom_up_all_classes.append(ancestor)
        i += 1
        class_an_des_list.append(i)
  for des_idx, descentant in enumerate(node_children): # select descendants
    if des_idx < max_depth_children:
      if ", " in descentant: # contains synonym, more than one, keep only one
        imagenet_bottom_up_all_classes.append(descentant.split(", ")[0])
        i += 1
        class_an_des_list.append(i)
      else:
        imagenet_bottom_up_all_classes.append(descentant)
        i += 1
        class_an_des_list.append(i)
  imagenet_bottom_up_mapping_index_list.append(class_an_des_list)

"""len(imagenet_bottom_up_mapping_index_list) = 1000, each element is the source node and its children"""

imagenet_classes_an_des, class_an_des_mapping, class_an_des_mapping_list = bottom_up_hierarchy(max_depth_descendant=10)
# len(imagenet_classes_an_des) = 1971. word
# len(class_an_des_mapping) = 1000. starting index of each class among 1971
# len(class_an_des_mapping_list) = 1000.

def compute_accuracy_details(logits, labels):
  top_probs, top_labels = jax.lax.top_k(logits, 5)
  top1 = 100 * jnp.mean(top_labels[:, 0] == labels)
  top5 = 100 * jnp.sum(top_labels == labels[:, None]) / labels.shape[0]
  return top1, top5, top_probs, top_labels

top1_prompt, top5_prompt, top5_prompt_probs, top5_prompt_labels = compute_accuracy_details(logits_prompteng, labels)
top1_name, top5_name, top5_name_probs, top5_name_labels = compute_accuracy_details(logits_name, labels)
print(f'Prompt Engineering: top1={top1_prompt_h:.2f}%, top5={top5_prompt:.2f}%')
print(f'Class Names: top1={top1_name:.2f}%, top5={top5_name:.2f}%')

# continuous uncertatinty value calculation
def continuous_uncertainty_estimation(prompts=clip.PROMPTS, name_prediction=top5_name_labels[:, 0]):
  """obtain continuous confidence score for each image"""

  # prompts decision
  prompts_slice_log = []
  for sample_id in range(len(prompts)): # for each different prompts
    prompteng_slice = zeroshot_classifier(clip.IMAGENET_CLASSES, [prompts[sample_id]]) #select one each time
    logits_slice = np.matmul(embeddings, prompteng_slice) # compute the logits
    probs_slice = np.array(jax.nn.softmax(logits_slice, axis=-1))
    preds_slice = jnp.argmax(probs_slice, axis=-1)
    prompts_slice_log.append(preds_slice)

  prompts_slice_log_array = np.array(prompts_slice_log) # (len(prompts), 50000)

  # compare with no-prompts and use consistency to get the confidence score
  consistency_with_nonpromp = prompts_slice_log_array == name_prediction
  consistency_with_nonpromp_log = sum(consistency_with_nonpromp, axis=0)
  confidence_sort_index = numpy.argsort(consistency_with_nonpromp_log) # form small to large

  return prompts_slice_log_array, consistency_with_nonpromp_log, confidence_sort_index

prompts_slice_log_array, consistency_with_nonpromp_log, confidence_sort_index = continuous_uncertainty_estimation()

# Calculate accuracy of the low confidence set
low_confident_set_size = 10000
low_confident_set_index = confidence_sort_index[:low_confident_set_size]
# Computed accuracy for the low confidence set
acc = np.mean(name_prediction[np.array(low_confident_set_index)] == labels[low_confident_set_index])
acc

"""## Hierarchy-CLIP

### Continuous Solution
"""

an_des_dict = an_des_dict_json

def zeroshot_classifier_hierarchy(classnames, templates, permute=False):
  zeroshot_weights = []
  permute_fn = permute_words if permute else lambda x: x
  for classname in classnames:
    texts = [permute_fn(template.format(classname.split('|')[0]) + template.format(classname.split('|')[1])) for template in templates]
    class_embeddings = encode_text(tokenize_fn(texts))
    class_embedding = class_embeddings.mean(0)
    class_embedding /= jnp.linalg.norm(class_embedding)
    zeroshot_weights.append(class_embedding)
  return jnp.stack(zeroshot_weights, axis=1)

def compute_accuracy_rerank(logits, ori_labels, labels):
  """based on logits, reorder the ori_top5_labels and compare with true labels"""

  top_probs, top_labels_index = jax.lax.top_k(logits, 1)
  top_labels = np.array([ele[top_labels_index[idx]] for idx, ele in enumerate(tqdm(ori_labels))])
  top1 = 100 * jnp.mean(top_labels[:, 0] == labels)
  return top1, top_probs, top_labels

# Bottom-up and top-down augmented inference

from collections import Counter

an_des_dict = an_des_dict_json # simplified wordnet
fill_empty_parent = ''
unstable_num = 1000
# stable set and unstable set
unstable_index = entropy_sort_index[:unstable_num] # number of rejected/unstable images
stable_index = np.array(list(set(range(len(labels))) - set(unstable_index.tolist()))) # stable, np.array

# 1. accuracy of the stable samples:
# Computed accuracy
acc_stable = np.mean(name_prediction[stable_index] == labels[stable_index])

acc_unstable_ori = np.mean(name_prediction[unstable_index] == labels[unstable_index])


# 2. accuracy of the unstable samples: use our hierarchy
# analysis top 5 in all different cases of baseline
baseline_top5 = top5_name_labels[unstable_index]
baseline_top5_probs = top5_name_probs[unstable_index]
diff_image_embeddings = embeddings[unstable_index]
unstable_id_gt_labels = labels[unstable_index]

unstable_case = top5_name_labels[unstable_index]   # original top-5 prediction
unstable_case_label = labels[unstable_index] # top-5 GT labels

# [Top-5 Ancestor collection] calculate Ancestor from simple wordnet for unstable cases
A_all = []
for example in unstable_case: # for each unstable image
  # ancestor of top 5
  top5_ancestor = []
  # top5_ancestor_raw = []
  for ind, top_k in enumerate(example):
    if an_des_dict[str(example[ind])]['ancestors']==[]: # no ancestor
      top5_ancestor.append(fill_empty_parent)
    else:
      top5_ancestor.append(an_des_dict[str(example[ind])]['ancestors'].copy()[0]) # avoid influence the root
  A_all.append(top5_ancestor)


# add ancestor to all candidate classes
logits_ori_all = []
for idx, unstable_case_ele in tqdm(enumerate(unstable_case)): # for each unstable image
  # original word embedding
  word_list_ori = np.array(clip.IMAGENET_CLASSES)[np.array(unstable_case_ele)].tolist()
  word_list_hierarchy = [] # raw words with only [bottom up]
  word_list_hierarchy_with_prompt = [] # words add prompt and ancestor [bottom up + top down]
  wordid_list_hierarchy_list = [] # details of index [[1,2], [3,4,5], [6], [7,8], [9]]
  wordid_list_hierarchy_mapping_list = [] # used for reduceat (index of first element for each top-5) [1,3,6,7,9]
  i = 0 # tracking each class/child class
  for word_id, word in enumerate(word_list_ori): # for each class in top 5
    tempid_list = [] # for specific word in top5
    wordid_list_hierarchy_mapping_list.append(i)
    for child_id in class_an_des_mapping_list[unstable_case_ele[word_id]]: # for each children or element [Bottom-up]
      raw_word = imagenet_classes_an_des[child_id-1]
      word_list_hierarchy.append(raw_word) # child_id start from 1 not 0
      word_list_hierarchy_with_prompt.append(raw_word + ' which is a kind of|' + A_all[idx][word_id]) # [Top-down]
      # word_list_hierarchy_with_prompt.append(raw_word + ' which is a kind of' + A_all[idx][word_id]) # [Top-down]
      tempid_list.append(i)
      i+=1
    wordid_list_hierarchy_list.append(tempid_list)
  word_list = word_list_hierarchy_with_prompt # could change

  # inference
  word_weights_ori = zeroshot_classifier_hierarchy(word_list, ['{}']) # (512, 5) # do not use prompt ensemble
  # word_weights_ori = zeroshot_classifier_hierarchy(word_list, clip.PROMPTS) # (512, 5)
  img_embedding = diff_image_embeddings[idx] # (1, 512) single image
  logits_ori_raw = np.matmul(img_embedding, word_weights_ori)
  logits_ori_raw = logits_ori_raw[np.newaxis, :]
  logits_ori = np.maximum.reduceat(logits_ori_raw, indices=wordid_list_hierarchy_mapping_list, axis=1)
  logits_ori_all.append(logits_ori)
logits_ori_all_array = np.array(logits_ori_all)[:, 0, :]

# Manually computed accuracy
unstable_top1_acc, unstable_top_probs, unstable_top_labels = compute_accuracy_rerank(logits_ori_all_array, unstable_case, unstable_case_label)


# compute overall accuracy
overall_acc = (acc_stable*100*len(stable_index) + unstable_top1_acc*len(unstable_index))/len(labels)

print("acc_stable:", acc_stable)
print("acc_unstable_ori:", acc_unstable_ori)
print("acc_unstable_with_Hierarchy-CLIP:", unstable_top1_acc)
print("overall_acc_with_Hierarchy-CLIP:", overall_acc)

"""### Descrete solution"""

# first half
prompteng_firsthalf = zeroshot_classifier(clip.IMAGENET_CLASSES, clip.PROMPTS[0:40]) #select one each time
logits_firsthalf = np.matmul(embeddings, prompteng_firsthalf) # compute the logits
probs_firsthalf = np.array(jax.nn.softmax(logits_firsthalf, axis=-1))
preds_firsthalf = jnp.argmax(probs_firsthalf, axis=-1)

# 2nd half
prompteng_2ndhalf = zeroshot_classifier(clip.IMAGENET_CLASSES, clip.PROMPTS[40:80]) #select one each time
logits_2ndhalf = np.matmul(embeddings, prompteng_2ndhalf) # compute the logits
probs_2ndhalf = np.array(jax.nn.softmax(logits_2ndhalf, axis=-1))
preds_2ndhalf = jnp.argmax(probs_2ndhalf, axis=-1)

name_prediction=top5_name_labels[:, 0]
prompt_prediction = top5_prompt_labels[:, 0]

prompts_log = np.stack((name_prediction, prompt_prediction, preds_firsthalf, preds_2ndhalf))

stable_log = np.array([int(np.all(prompts_log[:, i] == prompts_log[:, i][0])) for i in tqdm(range(prompts_log.shape[-1]))])

unstable_id = np.where(stable_log==0)

# discrete uncertatinty value calculation
def discrete_uncertainty_estimation(prompts=clip.PROMPTS, name_prediction=top5_name_labels[:, 0], prompt_prediction=top5_prompt_labels[:, 0]):
  """obtain discrete confidence score for each image"""

  # first half
  prompteng_firsthalf = zeroshot_classifier(clip.IMAGENET_CLASSES, prompts[0:int(len(prompts)/2)]) #first half [0:40]
  logits_firsthalf = np.matmul(embeddings, prompteng_firsthalf) # compute the logits
  probs_firsthalf = np.array(jax.nn.softmax(logits_firsthalf, axis=-1))
  preds_firsthalf = jnp.argmax(probs_firsthalf, axis=-1)

  # 2nd half
  prompteng_2ndhalf = zeroshot_classifier(clip.IMAGENET_CLASSES, prompts[int(len(prompts)/2):]) #second half [40:80]
  logits_2ndhalf = np.matmul(embeddings, prompteng_2ndhalf) # compute the logits
  probs_2ndhalf = np.array(jax.nn.softmax(logits_2ndhalf, axis=-1))
  preds_2ndhalf = jnp.argmax(probs_2ndhalf, axis=-1)

  # consider all four different kinds of prompts
  prompts_log = np.stack((name_prediction, prompt_prediction, preds_firsthalf, preds_2ndhalf))

  consistency_log = np.array([int(np.all(prompts_log[:, i] == prompts_log[:, i][0])) for i in tqdm(range(prompts_log.shape[-1]))])

  unstable_id = np.where(consistency_log==0)

  return consistency_log, unstable_id

consistency_log, unstable_id = discrete_uncertainty_estimation()
# Computed accuracy for the low confidence set len(unstable_id[0]) = 10308
acc = np.mean(name_prediction[np.array(unstable_id)] == labels[unstable_id])
acc