Part-based, explainable and editable bird classifier that allows users to define in text the bird
import json
import os
import re
import openai
from tqdm import tqdm
from utils import PROJECT_ROOT, load_descriptions
import tiktoken
import time
from ast import literal_eval

openai.api_key = os.getenv("OPENAI_API_KEY")


def num_tokens_from_messages(messages, model="gpt-3.5-turbo-0301"):
    """Returns the number of tokens used by a list of messages."""
    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")
    if model in ["gpt-3.5-turbo-0301", "gpt-4"]:  # note: future models may deviate from this
        num_tokens = 0
        for message in messages:
            num_tokens += 4  # every message follows <im_start>{role/name}\n{content}<im_end>\n
            for key, value in message.items():
                num_tokens += len(encoding.encode(value))
                if key == "name":  # if there's a name, the role is omitted
                    num_tokens += -1  # role is always required and always 1 token
        num_tokens += 2  # every reply is primed with <im_start>assistant
        return num_tokens
    else:
        raise NotImplementedError(f"""num_tokens_from_messages() is not presently implemented for model {model}.
    See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens.""")


def generate_prompt_sachit(category_name: str):
    # you can replace the examples with whatever you want; these were random and worked, could be improved
    return f"""Q: What are useful visual features for distinguishing a lemur in a photo?
A: There are several useful visual features to tell there is a lemur in a photo:
- four-limbed primate
- black, grey, white, brown, or red-brown
- wet and hairless nose with curved nostrils
- long tail
- large eyes
- furry bodies
- clawed hands and feet

Q: What are useful visual features for distinguishing a television in a photo?
A: There are several useful visual features to tell there is a television in a photo:
- electronic device
- black or grey
- a large, rectangular screen
- a stand or mount to support the screen
- one or more speakers
- a power cord
- input ports for connecting to other devices
- a remote control

Q: What are useful features for distinguishing a {category_name} in a photo?
A: There are several useful visual features to tell there is a {category_name} in a photo:
-
"""


def generate_descriptors(dataset, model):

    if dataset in ["cub", "cub_sachit"]:
        templated_descriptions, _ = load_descriptions(dataset_name="cub", prompt_type=0)
        class_names = list(templated_descriptions.keys())
    elif dataset in ["inaturalist", "inaturalist_sachit"]:
        with open("../data/class_names/inaturalist_classes.json", "r") as input_file:
            class_names = json.load(input_file)
    elif dataset in ["nabirds", "nabirds_sachit"]:
        with open("../data/class_names/nabirds_classes.json", "r") as input_file:
            class_names = json.load(input_file)
    elif dataset == "bird_soup":
        with open("../data/class_names/bird_soup_uncased_v2_classes.txt", "r") as input_file:
            class_names = [class_name.strip() for class_name in input_file.readlines()]
    elif dataset == "ebird":
        with open("../data/class_names/ebird_classes.txt", "r") as input_file:
            class_names = [class_name.strip() for class_name in input_file.readlines()]
    elif dataset == "bird_soup_scientific":
        with open("../data/class_names/cub_nabird_inat_sci.txt", "r") as input_file:
            class_names = [class_name.strip() for class_name in input_file.readlines()]

    if dataset in ["cub", "inaturalist", "nabirds", "bird_soup", "ebird", "bird_soup_extra", "bird_soup_scientific"]:
        content = "A bird has 12 parts: back, beak, belly, breast, crown, forehead, eyes, legs, wings, nape, tail and throat. " \
                  "Visually describe all parts of {} bird in a short phrase in bullet points using the format 'part: short phrase'"
    else:
        raise Exception(f"Dataset {dataset} is not supported")

    save_dir = f"{PROJECT_ROOT}/data/descriptors/chatgpt/raw/"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)

    raw_descriptors_fp = f"{save_dir}/descriptors_{dataset}.json"
    descriptors_fp = f"{save_dir.replace('raw/', '')}/descriptors_{dataset}.json"
    chat_results = {}
    all_messages = []

    if os.path.exists(raw_descriptors_fp):
        with open(raw_descriptors_fp, "r") as input_file:
            chat_results = json.load(input_file)

    loaded_class_names = list(chat_results.keys())

    # overwriten_class_names = ["tool kit", "mask", "necklace", "scarf", "website", "menu"]
    # overwriten_class_names = ["butterfly coquette", "cerulean cuckooshrike", "eye ringed flatbill", "magpie tanager", "red wattled lapwing"]
    overwriten_class_names = []

    index = 0
    with tqdm(total=len(class_names)) as pbar:
        while index < len(class_names):
            if class_names[index] in loaded_class_names and class_names[index] not in overwriten_class_names:
                index += 1
                pbar.update(1)
                continue

            try:
                if dataset in ["cub_sachit", "nabirds_sachit", "inaturalist_sachit"]:
                    content = generate_prompt_sachit(class_names[index])
                    temperature = 0.7
                else:
                    content = content.format(class_names[index], class_names[index])
                    temperature = 1.0

                completion = openai.ChatCompletion.create(
                    model=model,
                    messages=[
                        {"role": "user",
                         "content": content}
                    ],
                    temperature=temperature)
            except Exception as e:
                print(e)
                time.sleep(60)
                continue

            chat_results[class_names[index]] = completion.choices[0].message["content"]
            all_messages.append(completion.choices[0].message)

            num_tokens = num_tokens_from_messages([completion.choices[0].message], model)
            print(f"Number of tokens: {num_tokens}\n")

            with open(raw_descriptors_fp, "w") as output_file:
                json.dump(chat_results, output_file, indent=4)

            index += 1
            pbar.update(1)

    if not os.path.exists(descriptors_fp):
        with open(raw_descriptors_fp, "r") as input_file:
            chat_results = json.load(input_file)

        for class_name, chat_result in chat_results.items():
            descriptors = []
            for descriptor in chat_result.strip().replace("- ", "").split("\n"):
                if dataset not in ["cub_sachit", "nabirds_sachit", "inaturalist_sachit"]:
                    descriptor = re.sub(r"^\W+", "", descriptor.strip().lower())
                    descriptor = re.sub(r"\W+$", "", descriptor)

                if (descriptor and ":" in descriptor) or dataset in ["cub_sachit", "nabirds_sachit", "inaturalist_sachit"]:
                    descriptors.append(descriptor)
            chat_results[class_name] = descriptors

        with open(descriptors_fp, "w") as output_file:
            json.dump(chat_results, output_file, indent=4)
    else:
        with open(descriptors_fp, "r") as input_file:
            chat_results = json.load(input_file)

    total_num_tokens = 0
    for message in all_messages:
        total_num_tokens += num_tokens_from_messages([message], model)
    print(f"Total number of tokens: {total_num_tokens}\n")


def correct_cub_descriptors(model):

    with open(f"{PROJECT_ROOT}/data/descriptors/chatgpt/cub_allaboutbird_desc.json", "r") as input_file:
        cub_correct_descriptors = json.load(input_file)
        class_names = list(cub_correct_descriptors.keys())

    with open(f"{PROJECT_ROOT}/data/descriptors/chatgpt/descriptors_bird_soup_v21.json", "r") as input_file:
        cub_descriptors = json.load(input_file)

    # with open(f"{PROJECT_ROOT}/data/descriptors/chatgpt/descriptors_cub_sachit.json", "r") as input_file:
    #     cub_descriptors = json.load(input_file)
    #     cub_descriptors = {k.lower().replace("'s", "").replace("-", " ").replace("_", " ").replace("forsters tern", "forster tern"): v for k, v in cub_descriptors.items()}

    content = """Given the description of the bird "{}" as follows:

{}

Please summarize the visual features. Then check the items below one by one, if the item does not match or cannot be inferred from the description, correct it. If no information is provided, use your own knowledge.

{}

Please return with the same format. Note that if not specified, simply keep the original description.
    """

    save_dir = f"{PROJECT_ROOT}/data/descriptors/chatgpt/raw"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)

    raw_descriptors_fp = f"{save_dir}/descriptors_cub_auto_correct.json"                    # _sachit
    descriptors_fp = f"{save_dir.replace('/raw', '')}/descriptors_cub_auto_correct.json"    # _sachit
    chat_results = {}
    all_messages = []

    if os.path.exists(raw_descriptors_fp):
        with open(raw_descriptors_fp, "r") as input_file:
            chat_results = json.load(input_file)

    loaded_class_names = list(chat_results.keys())

    overwriten_class_names = []
    index = 0
    questions = {}

    for idx, class_name in enumerate(class_names):
        old_descs = "[\n"
        for item in cub_descriptors[class_name]:
            old_descs += "\t\"" + item + "\",\n"
        old_descs += "]"

        description = content.format(class_name, cub_correct_descriptors[class_name], old_descs)
        questions[class_name] = description

    with open("questions.json", "w") as output_file:                                        # _sachit
        json.dump(questions, output_file, indent=4)

    # for idx, class_name in tqdm(enumerate(class_names)):
    with tqdm(total=len(class_names)) as pbar:
        while index < len(class_names):
            if class_names[index] in loaded_class_names and class_names[index] not in overwriten_class_names:
                index += 1
                pbar.update(1)
                continue

            try:
                old_descs = "[\n"
                for item in cub_descriptors[class_names[index]]:
                    old_descs += "\t\"" + item + "\",\n"
                old_descs += "]"

                description = content.format(class_names[index], cub_correct_descriptors[class_names[index]], old_descs)
                temperature = 1.0
                if index < 5:
                    print(description)

                completion = openai.ChatCompletion.create(
                    model=model,
                    messages=[
                        {"role": "user",
                         "content": description}
                    ],
                    temperature=temperature)
            except Exception as e:
                print(e)
                time.sleep(60)
                continue

            chat_results[class_names[index]] = completion.choices[0].message["content"]
            all_messages.append(completion.choices[0].message)

            num_tokens = num_tokens_from_messages([completion.choices[0].message], model)
            print(f"Number of tokens: {num_tokens}\n")

            with open(raw_descriptors_fp, "w") as output_file:
                json.dump(chat_results, output_file, indent=4)

            index += 1
            pbar.update(1)

    if not os.path.exists(descriptors_fp):
        with open(raw_descriptors_fp, "r") as input_file:
            chat_results = json.load(input_file)

        for class_name, chat_result in chat_results.items():
            chat_results[class_name] = literal_eval(chat_result)

        with open(descriptors_fp, "w") as output_file:
            json.dump(chat_results, output_file, indent=4)
    else:
        with open(descriptors_fp, "r") as input_file:
            chat_results = json.load(input_file)

    total_num_tokens = 0
    for message in all_messages:
        total_num_tokens += num_tokens_from_messages([message], model)
    print(f"Total number of tokens: {total_num_tokens}\n")


def test_chatgpt(model, content):
    print(content)
    print(f"Response from ChatGPT version {model}:")

    completion = openai.ChatCompletion.create(
        model=model,
        messages=[
            {"role": "user",
             "content": content}
        ])

    print(completion.choices[0].message["content"])

    num_tokens = num_tokens_from_messages([completion.choices[0].message], model)
    print(f"Number of tokens: {num_tokens}\n")

    return num_tokens


if __name__ == '__main__':
    model = "gpt-4"  # gpt-3.5-turbo, gpt-4

    generate_descriptors(dataset="cub", model=model)
    # generate_descriptors(dataset="nabirds", model=model)
    # generate_descriptors(dataset="bird_soup_extra", model=model)
    # generate_descriptors(dataset="bird_soup_scientific", model=model)
    # correct_cub_descriptors(model=model)

src/chatgpt.py
anguyen8/peeb
from pathlib import Path
import torch
import os


def get_project_root() -> Path:
    """Returns project root folder."""
    return Path(__file__).parent.parent


PROJECT_ROOT = get_project_root()

device = "cuda" if torch.cuda.is_available() else "cpu"

# ------------------------------------------
# TODO: Change this to your own directory
# ------------------------------------------
# DATASET_DIR = ''
# PRECOMPUTED_DIR = ''

DATASET_DIR = '/home/lab/datasets'
PRECOMPUTED_DIR = '/home/lab/xclip/part_boxes/owlvit-large-patch14_cub-12-parts'
# ------------------------------------------

CUB_DIR = os.path.join(DATASET_DIR, 'cub/CUB_200_2011')
NABIRDS_DIR = os.path.join(DATASET_DIR, 'nabirds')
BIRD_SOUP_DIR = os.path.join(DATASET_DIR, 'bird_soup')

DS_CUB = "cub"
DS_NABIRDS = "nabirds"
DS_BIRD_SOUP = "bird_soup"

# For updating boxes logits with DATASET_WRAPPER
# BIRD_SOUP_META_PATH_ALL = f"{BIRD_SOUP_DIR}/metadata/bird_soup_uncased.h5"    # 343K
BIRD_SOUP_META_PATH_ALL = os.path.join(BIRD_SOUP_DIR, "metadata/bird_soup_uncased_v2.h5")   # 398K --> 440K

# BIRD_SOUP_META_PATH_TEST  = f"{BIRD_SOUP_DIR}/metadata/v2/pretrain_exclude_cub/cub_classes_all.h5"
# BIRD_SOUP_META_PATH_TEST  = f"{BIRD_SOUP_DIR}/metadata/cub_test_reindexed.h5"
# BIRD_SOUP_META_PATH_TEST  = f"{BIRD_SOUP_DIR}/metadata/nabirds_test_reindexed.h5"
# BIRD_SOUP_META_PATH_TEST  = f"{BIRD_SOUP_DIR}/metadata/inaturalist_filtered.h5"

# For BirdSoup-v2.1 (Level 1)
BIRD_SOUP_META_PATH_TRAIN = f"{BIRD_SOUP_DIR}/metadata/level_1_exclude_cub_nabirds_inat/train_keep_child_a100_reindexed.h5"
BIRD_SOUP_META_PATH_VAL   = f"{BIRD_SOUP_DIR}/metadata/level_1_exclude_cub_nabirds_inat/val_keep_child_a100_reindexed.h5"
BIRD_SOUP_META_PATH_TEST  = f"{BIRD_SOUP_DIR}/metadata/level_1_exclude_cub_nabirds_inat/test_cub_reindexed.h5"

# For BirdSoup-v2.1 (Level 3)
# CUB
# BIRD_SOUP_META_PATH_TRAIN = f"{BIRD_SOUP_DIR}/metadata/level_3_exclude_cub/train_keep_child_a100_reindexed.h5"
# BIRD_SOUP_META_PATH_VAL   = f"{BIRD_SOUP_DIR}/metadata/level_3_exclude_cub/val_keep_child_a100_reindexed.h5"
# BIRD_SOUP_META_PATH_TEST  = f"{BIRD_SOUP_DIR}/metadata/level_3_exclude_cub/clore_split/test_unseen_reindexed.h5"

# NABIRDS
# BIRD_SOUP_META_PATH_TRAIN = f"{BIRD_SOUP_DIR}/metadata/level_3_exclude_nabirds/train_keep_child_a100_reindexed.h5"
# BIRD_SOUP_META_PATH_VAL   = f"{BIRD_SOUP_DIR}/metadata/level_3_exclude_nabirds/val_keep_child_a100_reindexed.h5"
# BIRD_SOUP_META_PATH_TEST  = f"{BIRD_SOUP_DIR}/metadata/level_1_exclude_cub_nabirds_inat/test_nabirds_reindexed.h5"

# INATURALIST
# BIRD_SOUP_META_PATH_TRAIN = f"{BIRD_SOUP_DIR}/metadata/level_3_exclude_inaturalist/train_keep_child_a100_reindexed.h5"
# BIRD_SOUP_META_PATH_VAL   = f"{BIRD_SOUP_DIR}/metadata/level_3_exclude_inaturalist/val_keep_child_a100_reindexed.h5"
# BIRD_SOUP_META_PATH_TEST  = f"{BIRD_SOUP_DIR}/metadata/level_1_exclude_cub_nabirds_inat/test_inat_reindexed.h5"

thousand_k_to_200 = {0: -1, 1: -1, 2: -1, 3: -1, 4: -1, 5: -1, 6: 0, 7: -1, 8: -1, 9: -1, 10: -1, 11: 1, 12: -1,
                     13: 2, 14: -1, 15: 3, 16: -1, 17: 4, 18: -1, 19: -1, 20: -1, 21: -1, 22: 5, 23: 6, 24: -1,
                     25: -1, 26: -1, 27: 7, 28: -1, 29: -1, 30: 8, 31: -1, 32: -1, 33: -1, 34: -1, 35: -1, 36: -1,
                     37: 9, 38: -1, 39: 10, 40: -1, 41: -1, 42: 11, 43: -1, 44: -1, 45: -1, 46: -1, 47: 12, 48: -1,
                     49: -1, 50: 13, 51: -1, 52: -1, 53: -1, 54: -1, 55: -1, 56: -1, 57: 14, 58: -1, 59: -1, 60: -1,
                     61: -1, 62: -1, 63: -1, 64: -1, 65: -1, 66: -1, 67: -1, 68: -1, 69: -1, 70: 15, 71: 16, 72: -1,
                     73: -1, 74: -1, 75: -1, 76: 17, 77: -1, 78: -1, 79: 18, 80: -1, 81: -1, 82: -1, 83: -1, 84: -1,
                     85: -1, 86: -1, 87: -1, 88: -1, 89: 19, 90: 20, 91: -1, 92: -1, 93: -1, 94: 21, 95: -1, 96: 22,
                     97: 23, 98: -1, 99: 24, 100: -1, 101: -1, 102: -1, 103: -1, 104: -1, 105: 25, 106: -1, 107: 26,
                     108: 27, 109: -1, 110: 28, 111: -1, 112: -1, 113: 29, 114: -1, 115: -1, 116: -1, 117: -1,
                     118: -1, 119: -1, 120: -1, 121: -1, 122: -1, 123: -1, 124: 30, 125: 31, 126: -1, 127: -1,
                     128: -1, 129: -1, 130: 32, 131: -1, 132: 33, 133: -1, 134: -1, 135: -1, 136: -1, 137: -1,
                     138: -1, 139: -1, 140: -1, 141: -1, 142: -1, 143: 34, 144: 35, 145: -1, 146: -1, 147: -1,
                     148: -1, 149: -1, 150: 36, 151: 37, 152: -1, 153: -1, 154: -1, 155: -1, 156: -1, 157: -1,
                     158: -1, 159: -1, 160: -1, 161: -1, 162: -1, 163: -1, 164: -1, 165: -1, 166: -1, 167: -1,
                     168: -1, 169: -1, 170: -1, 171: -1, 172: -1, 173: -1, 174: -1, 175: -1, 176: -1, 177: -1,
                     178: -1, 179: -1, 180: -1, 181: -1, 182: -1, 183: -1, 184: -1, 185: -1, 186: -1, 187: -1,
                     188: -1, 189: -1, 190: -1, 191: -1, 192: -1, 193: -1, 194: -1, 195: -1, 196: -1, 197: -1,
                     198: -1, 199: -1, 200: -1, 201: -1, 202: -1, 203: -1, 204: -1, 205: -1, 206: -1, 207: 38,
                     208: -1, 209: -1, 210: -1, 211: -1, 212: -1, 213: -1, 214: -1, 215: -1, 216: -1, 217: -1,
                     218: -1, 219: -1, 220: -1, 221: -1, 222: -1, 223: -1, 224: -1, 225: -1, 226: -1, 227: -1,
                     228: -1, 229: -1, 230: -1, 231: -1, 232: -1, 233: -1, 234: 39, 235: 40, 236: -1, 237: -1,
                     238: -1, 239: -1, 240: -1, 241: -1, 242: -1, 243: -1, 244: -1, 245: -1, 246: -1, 247: -1,
                     248: -1, 249: -1, 250: -1, 251: -1, 252: -1, 253: -1, 254: 41, 255: -1, 256: -1, 257: -1,
                     258: -1, 259: -1, 260: -1, 261: -1, 262: -1, 263: -1, 264: -1, 265: -1, 266: -1, 267: -1,
                     268: -1, 269: -1, 270: -1, 271: -1, 272: -1, 273: -1, 274: -1, 275: -1, 276: -1, 277: 42,
                     278: -1, 279: -1, 280: -1, 281: -1, 282: -1, 283: 43, 284: -1, 285: -1, 286: -1, 287: 44,
                     288: -1, 289: -1, 290: -1, 291: 45, 292: -1, 293: -1, 294: -1, 295: 46, 296: -1,  297: -1,
                     298: 47, 299: -1, 300: -1, 301: 48, 302: -1, 303: -1, 304: -1, 305: -1, 306: 49, 307: 50,
                     308: 51, 309: 52, 310: 53, 311: 54, 312: -1, 313: 55, 314: 56, 315: 57, 316: -1, 317: 58,
                     318: -1, 319: 59, 320: -1, 321: -1, 322: -1, 323: 60, 324: 61, 325: -1, 326: 62, 327: 63,
                     328: -1, 329: -1, 330: 64, 331: -1, 332: -1, 333: -1, 334: 65, 335: 66, 336: 67, 337: -1,
                     338: -1, 339: -1, 340: -1, 341: -1, 342: -1, 343: -1, 344: -1, 345: -1, 346: -1, 347: 68,
                     348: -1, 349: -1, 350: -1, 351: -1, 352: -1, 353: -1, 354: -1, 355: -1, 356: -1, 357: -1,
                     358: -1, 359: -1, 360: -1, 361: 69, 362: -1, 363: 70, 364: -1, 365: -1, 366: -1, 367: -1,
                     368: -1, 369: -1, 370: -1, 371: -1, 372: 71, 373: -1, 374: -1, 375: -1, 376: -1, 377: -1,
                     378: 72, 379: -1, 380: -1, 381: -1, 382: -1, 383: -1, 384: -1, 385: -1, 386: 73, 387: -1,
                     388: -1, 389: -1, 390: -1, 391: -1, 392: -1, 393: -1, 394: -1, 395: -1, 396: -1, 397: 74,
                     398: -1, 399: -1, 400: 75, 401: 76, 402: 77, 403: -1, 404: 78, 405: -1, 406: -1, 407: 79,
                     408: -1, 409: -1, 410: -1, 411: 80, 412: -1, 413: -1, 414: -1, 415: -1, 416: 81, 417: 82,
                     418: -1, 419: -1, 420: 83, 421: -1, 422: -1, 423: -1, 424: -1, 425: 84, 426: -1, 427: -1,
                     428: 85, 429: -1, 430: 86, 431: -1, 432: -1, 433: -1, 434: -1, 435: -1, 436: -1, 437: 87,
                     438: 88, 439: -1, 440: -1, 441: -1, 442: -1, 443: -1, 444: -1, 445: 89, 446: -1, 447: -1,
                     448: -1, 449: -1, 450: -1, 451: -1, 452: -1, 453: -1, 454: -1, 455: -1, 456: 90, 457: 91,
                     458: -1, 459: -1, 460: -1, 461: 92, 462: 93, 463: -1, 464: -1, 465: -1, 466: -1, 467: -1,
                     468: -1, 469: -1, 470: 94, 471: -1, 472: 95, 473: -1, 474: -1, 475: -1, 476: -1, 477: -1,
                     478: -1, 479: -1, 480: -1, 481: -1, 482: -1, 483: 96, 484: -1, 485: -1, 486: 97, 487: -1,
                     488: 98, 489: -1, 490: -1, 491: -1, 492: 99, 493: -1, 494: -1, 495: -1, 496: 100, 497: -1,
                     498: -1, 499: -1, 500: -1, 501: -1, 502: -1, 503: -1, 504: -1, 505: -1, 506: -1, 507: -1,
                     508: -1, 509: -1, 510: -1, 511: -1, 512: -1, 513: -1, 514: 101, 515: -1, 516: 102, 517: -1,
                     518: -1, 519: -1, 520: -1, 521: -1, 522: -1, 523: -1, 524: -1, 525: -1, 526: -1, 527: -1,
                     528: 103, 529: -1, 530: 104, 531: -1, 532: -1, 533: -1, 534: -1, 535: -1, 536: -1, 537: -1,
                     538: -1, 539: 105, 540: -1, 541: -1, 542: 106, 543: 107, 544: -1, 545: -1, 546: -1, 547: -1,
                     548: -1, 549: 108, 550: -1, 551: -1, 552: 109, 553: -1, 554: -1, 555: -1, 556: -1, 557: 110,
                     558: -1, 559: -1, 560: -1, 561: 111, 562: 112, 563: -1, 564: -1, 565: -1, 566: -1, 567: -1,
                     568: -1, 569: 113, 570: -1, 571: -1, 572: 114, 573: 115, 574: -1, 575: 116, 576: -1, 577: -1,
                     578: -1, 579: 117, 580: -1, 581: -1, 582: -1, 583: -1, 584: -1, 585: -1, 586: -1, 587: -1,
                     588: -1, 589: 118, 590: -1, 591: -1, 592: -1, 593: -1, 594: -1, 595: -1, 596: -1, 597: -1,
                     598: -1, 599: -1, 600: -1, 601: -1, 602: -1, 603: -1, 604: -1, 605: -1, 606: 119, 607: 120,
                     608: -1, 609: 121, 610: -1, 611: -1, 612: -1, 613: -1, 614: 122, 615: -1, 616: -1, 617: -1,
                     618: -1, 619: -1, 620: -1, 621: -1, 622: -1, 623: -1, 624: -1, 625: -1, 626: 123, 627: 124,
                     628: -1, 629: -1, 630: -1, 631: -1, 632: -1, 633: -1, 634: -1, 635: -1, 636: -1, 637: -1,
                     638: -1, 639: -1, 640: 125, 641: 126, 642: 127, 643: 128, 644: -1, 645: -1, 646: -1, 647: -1,
                     648: -1, 649: -1, 650: -1, 651: -1, 652: -1, 653: -1, 654: -1, 655: -1, 656: -1, 657: -1,
                     658: 129, 659: -1, 660: -1, 661: -1, 662: -1, 663: -1, 664: -1, 665: -1, 666: -1, 667: -1,
                     668: 130, 669: -1, 670: -1, 671: -1, 672: -1, 673: -1, 674: -1, 675: -1, 676: -1, 677: 131,
                     678: -1, 679: -1, 680: -1, 681: -1, 682: 132, 683: -1, 684: 133, 685: -1, 686: -1, 687: 134,
                     688: -1, 689: -1, 690: -1, 691: -1, 692: -1, 693: -1, 694: -1, 695: -1, 696: -1, 697: -1,
                     698: -1, 699: -1, 700: -1, 701: 135, 702: -1, 703: -1, 704: 136, 705: -1, 706: -1, 707: -1,
                     708: -1, 709: -1, 710: -1, 711: -1, 712: -1, 713: -1, 714: -1, 715: -1, 716: -1, 717: -1,
                     718: -1, 719: 137, 720: -1, 721: -1, 722: -1, 723: -1, 724: -1, 725: -1, 726: -1, 727: -1,
                     728: -1, 729: -1, 730: -1, 731: -1, 732: -1, 733: -1, 734: -1, 735: -1, 736: 138, 737: -1,
                     738: -1, 739: -1, 740: -1, 741: -1, 742: -1, 743: -1, 744: -1, 745: -1, 746: 139, 747: -1,
                     748: -1, 749: 140, 750: -1, 751: -1, 752: 141, 753: -1, 754: -1, 755: -1, 756: -1, 757: -1,
                     758: 142, 759: -1, 760: -1, 761: -1, 762: -1, 763: 143, 764: -1, 765: 144, 766: -1, 767: -1,
                     768: 145, 769: -1, 770: -1, 771: -1, 772: -1, 773: 146, 774: 147, 775: -1, 776: 148, 777: -1,
                     778: -1, 779: 149, 780: 150, 781: -1, 782: -1, 783: -1, 784: -1, 785: -1, 786: 151, 787: -1,
                     788: -1, 789: -1, 790: -1, 791: -1, 792: 152, 793: -1, 794: -1, 795: -1, 796: -1, 797: 153,
                     798: -1, 799: -1, 800: -1, 801: -1, 802: 154, 803: 155, 804: 156, 805: -1, 806: -1, 807: -1,
                     808: -1, 809: -1, 810: -1, 811: -1, 812: -1, 813: 157, 814: -1, 815: 158, 816: -1, 817: -1,
                     818: -1, 819: -1, 820: 159, 821: -1, 822: -1, 823: 160, 824: -1, 825: -1, 826: -1, 827: -1,
                     828: -1, 829: -1, 830: -1, 831: 161, 832: -1, 833: 162, 834: -1, 835: 163, 836: -1, 837: -1,
                     838: -1, 839: 164, 840: -1, 841: -1, 842: -1, 843: -1, 844: -1, 845: 165, 846: -1, 847: 166,
                     848: -1, 849: -1, 850: 167, 851: -1, 852: -1, 853: -1, 854: -1, 855: -1, 856: -1, 857: -1,
                     858: -1, 859: 168, 860: -1, 861: -1, 862: 169, 863: -1, 864: -1, 865: -1, 866: -1, 867: -1,
                     868: -1, 869: -1, 870: 170, 871: -1, 872: -1, 873: -1, 874: -1, 875: -1, 876: -1, 877: -1,
                     878: -1, 879: 171, 880: 172, 881: -1, 882: -1, 883: -1, 884: -1, 885: -1, 886: -1, 887: -1,
                     888: 173, 889: -1, 890: 174, 891: -1, 892: -1, 893: -1, 894: -1, 895: -1, 896: -1, 897: 175,
                     898: -1, 899: -1, 900: 176, 901: -1, 902: -1, 903: -1, 904: -1, 905: -1, 906: -1, 907: 177,
                     908: -1, 909: -1, 910: -1, 911: -1, 912: -1, 913: 178, 914: -1, 915: -1, 916: -1, 917: -1,
                     918: -1, 919: -1, 920: -1, 921: -1, 922: -1, 923: -1, 924: 179, 925: -1, 926: -1, 927: -1,
                     928: -1, 929: -1, 930: -1, 931: -1, 932: 180, 933: 181, 934: 182, 935: -1, 936: -1, 937: 183,
                     938: -1, 939: -1, 940: -1, 941: -1, 942: -1, 943: 184, 944: -1, 945: 185, 946: -1, 947: 186,
                     948: -1, 949: -1, 950: -1, 951: 187, 952: -1, 953: -1, 954: 188, 955: -1, 956: 189,957: 190,
                     958: -1, 959: 191, 960: -1, 961: -1, 962: -1, 963: -1, 964: -1, 965: -1, 966: -1,967: -1,
                     968: -1, 969: -1, 970: -1, 971: 192, 972: 193, 973: -1, 974: -1, 975: -1, 976: -1, 977: -1,
                     978: -1, 979: -1, 980: 194, 981: 195, 982: -1, 983: -1, 984: 196, 985: -1, 986: 197, 987: 198,
                     988: 199, 989: -1, 990: -1, 991: -1, 992: -1, 993: -1, 994: -1, 995: -1, 996: -1, 997: -1,
                     998: -1, 999: -1}

indices_in_1k = [k for k in thousand_k_to_200 if thousand_k_to_200[k] != -1]

src/configs.py
anguyen8/peeb
from .dataset_wrapper import DatasetWrapper
from .cub import CUB
from .nabirds import NABirdsDataset
from .bird_soup import BirdSoup
from .augmentation import TrivialAugmentation

src/data_loader/__init__.py
anguyen8/peeb
from typing import Dict, List, Optional, Tuple
from torch import Tensor
import torchvision.transforms.functional as F
from torchvision.transforms import InterpolationMode
from torchvision.transforms.autoaugment import _apply_op
import torch


class TrivialAugmentation(torch.nn.Module):
    r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in
    `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" <https://arxiv.org/abs/2103.10158>`_.
    If the image is torch Tensor, it should be of type torch.uint8, and it is expected
    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, it is expected to be in mode "L" or "RGB".

    Args:
        num_magnitude_bins (int): The number of different magnitude values.
        interpolation (InterpolationMode): Desired interpolation enum defined by
            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
        fill (sequence or number, optional): Pixel fill value for the area outside the transformed
            image. If given a number, the value is used for all bands respectively.
    """
    available_ops = ['Identity', 'ShearX', 'ShearY', 'TranslateX', 'TranslateY', 'Rotate', 'Brightness', 'Color', 'Contrast', 'Sharpness', 'Posterize', 'Solarize', 'AutoContrast', 'Equalize']

    def __init__(
        self,
        num_magnitude_bins: int = 31,
        interpolation: InterpolationMode = InterpolationMode.NEAREST,
        fill: Optional[List[float]] = None,
        operations: Optional[List[str]] = None,
    ) -> None:
        super().__init__()
        self.num_magnitude_bins = num_magnitude_bins
        self.interpolation = interpolation
        self.fill = fill
        self.op_meta = self._augmentation_space(num_bins=num_magnitude_bins, operations=operations)
        # if operations is not None, check if the operations are valid
        if operations is not None:
            for op in operations:
                assert op in self.available_ops, f"Operation {op} is not supported. Please choose from {self.available_ops}"

    def _augmentation_space(self, num_bins: int, operations: List[str]) -> Dict[str, Tuple[Tensor, bool]]:
        all_op = {
            # op_name: (magnitudes, signed)
            "Identity": (torch.tensor(0.0), False),
            "ShearX": (torch.linspace(0.0, 0.99, num_bins), True),
            "ShearY": (torch.linspace(0.0, 0.99, num_bins), True),
            "TranslateX": (torch.linspace(0.0, 32.0, num_bins), True),
            "TranslateY": (torch.linspace(0.0, 32.0, num_bins), True),
            "Rotate": (torch.linspace(0.0, 135.0, num_bins), True),
            "Brightness": (torch.linspace(0.0, 0.99, num_bins), True),
            "Color": (torch.linspace(0.0, 0.99, num_bins), True),
            "Contrast": (torch.linspace(0.0, 0.99, num_bins), True),
            "Sharpness": (torch.linspace(0.0, 0.99, num_bins), True),
            "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 6)).round().int(), False),
            "Solarize": (torch.linspace(255.0, 0.0, num_bins), False),
            "AutoContrast": (torch.tensor(0.0), False),
            "Equalize": (torch.tensor(0.0), False),
        }
        if operations is None:
            return all_op
        else:
            return {k: v for k, v in all_op.items() if k in operations}

    def forward(self, img: Tensor) -> Tensor:
        """
            img (PIL Image or Tensor): Image to be transformed.

        Returns:
            PIL Image or Tensor: Transformed image.
        """
        fill = self.fill
        channels, height, width = F.get_dimensions(img)
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * channels
            elif fill is not None:
                fill = [float(f) for f in fill]

        # op_meta = self._augmentation_space(self.num_magnitude_bins) # moved to __init__
        op_index = int(torch.randint(len(self.op_meta), (1,)).item())
        op_name = list(self.op_meta.keys())[op_index]
        magnitudes, signed = self.op_meta[op_name]

        # PEIJIE: fix magnitudes for float images input in Solarize opertation
        magnitudes = magnitudes.float() / 255.0 if op_name == 'Solarize' else magnitudes
        magnitude = (
            float(magnitudes[torch.randint(len(magnitudes), (1,), dtype=torch.long)].item())
            if magnitudes.ndim > 0
            else 0.0
        )
        if signed and torch.randint(2, (1,)):
            magnitude *= -1.0

        return _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill)

    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f"num_magnitude_bins={self.num_magnitude_bins}"
            f", interpolation={self.interpolation}"
            f", fill={self.fill}"
            f")"
        )
        return s



src/data_loader/augmentation.py
anguyen8/peeb
import math
import os
import json
import random

import torch
from PIL import Image
import pandas as pd
from torchvision import transforms
from torchvision.transforms import Compose
from torch.utils.data import Dataset, Subset, Sampler
from torchvision.datasets.folder import default_loader
from src.data_loader.augmentation import TrivialAugmentation
from torch.utils.data.distributed import DistributedSampler
from collections import defaultdict, deque
import copy


class BirdSoup(Dataset):

    def __init__(self, root: str, meta_path: str, box_dir: str = None, transform: Compose = None, use_meta_dir: bool = False,
                 train=True, subset: list[str] = None, return_path: bool = False):

        self.root_dir = root

        self.transform = transform
        self.train = train

        self.meta_df = pd.read_hdf(meta_path)

        self.use_meta_dir = use_meta_dir
        self.box_dir = box_dir
        self.subset = subset        # list of dataset names

        self.return_path = return_path

        if root is None and not use_meta_dir:
            raise ValueError("root and use_meta_dir cannot be both None")

        self._load_meta()

        # Data augmentation
        self.tri_aug = TrivialAugmentation(operations=[
            # 'ShearX',       # view point rotation
            # 'ShearY',
            # 'TranslateX',   # shifting
            # 'TranslateY',
            # 'Rotate',       # image rotation
            'Brightness',
            'Color',
            'Contrast',
            'Sharpness',
            # 'Posterize',  # quantize the color in 2^n, does not support float image
            # 'Solarize',   # invert the color below a threshold
            'AutoContrast',
            # 'Equalize'    # histogram equalization (the color will change), does not support float image
        ])
    
    def _load_meta(self,):
        data_sources = []
        if 'data_source' in self.meta_df.columns:
            data_sources = self.meta_df['data_source'].unique().tolist()

        if self.subset is not None:
            # TODO: CHECK RE_INDEXING OF CLASS_ID FOR SUB_CLASSES
            if "all" not in self.subset and set(self.subset).issubset(set(data_sources)):
                self.meta_df = self.meta_df[self.meta_df['data_source'].isin(self.subset)]

                # SORT dataframe by class_id
                self.meta_df = self.meta_df.sort_values(by=['class_id'])

                # and RE-INDEX the class_ids
                self.meta_df['class_id'] = self.meta_df['class_id'].astype('category').cat.codes

        self.targets = self.meta_df['class_id'].values.tolist()
        if self.use_meta_dir:
            image_paths = self.meta_df['abs_path'].values.tolist()
            self.samples = list(zip(image_paths, self.targets))
        else:
            self.meta_df['new_image_path'] = self.meta_df['new_image_name'].apply(lambda x: os.path.join(self.root_dir, "images", x))
            self.samples = list(zip(self.meta_df['new_image_path'].values.tolist(), self.targets))

        self.birdsoup_image_name = self.meta_df['new_image_name'].values.tolist()

        # Note: The unique() function returns the unique elements of the dataframe in top-bottom order.
        #      The order of the unique elements in the returned array generally preserves the original ordering.
        self.idx2class = dict(zip(self.meta_df['class_id'].unique(), self.meta_df['class_name'].unique()))
        self.class2idx = {v: k for k, v in self.idx2class.items()}

        if len(list(self.meta_df['class_id'].unique())) != len(list(self.meta_df['class_name'].unique())):
            print("WARNING: class_id and class_name are not 1-1 mapping")
        else:
            # sanity check: correctness of this mapping
            for k, v in self.idx2class.items():
                assert self.meta_df[self.meta_df['class_id'] == k]['class_name'].unique()[0] == v

        self.classes = list(self.class2idx.keys())
        self.imgs = self.samples

    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        image_path, target = self.samples[idx]
        sample = default_loader(image_path)
        soup_name = self.birdsoup_image_name[idx]
        size = torch.tensor(sample.size[::-1])      # (h, w)

        if self.train:
            sample = self.tri_aug(sample)

        if self.transform is not None:
            sample = self.transform(images=sample, return_tensor='pt')

        if self.return_path:
            return sample, target, image_path
        
        return sample, target, size, image_path, soup_name



src/data_loader/bird_soup.py
anguyen8/peeb
import os
import torch
from torchvision import datasets
from torchvision import transforms
from torchvision.transforms import Resize, InterpolationMode, ToTensor, Compose
from src.data_loader.augmentation import TrivialAugmentation

import random


class CUB(datasets.ImageFolder):
    """
    Wrapper for the CUB-200-2011 dataset.
    Method DatasetBirds.__getitem__() returns tuple of image and its corresponding label.
    Dataset per https://github.com/slipnitskaya/caltech-birds-advanced-classification
    """

    def __init__(self,
                 root,
                 transform=None,
                 target_transform=None,
                 loader=datasets.folder.default_loader,
                 is_valid_file=None,
                 train=True,
                 image_type=None,
                 zeroshot_split=False,
                 return_path=False):

        self.root_dir = root
        image_type = "images" if image_type is None else f"images_{image_type}"
        img_root = os.path.join(root, image_type)

        super(CUB, self).__init__(
            root=img_root,
            transform=None,
            target_transform=None,
            loader=loader,
            is_valid_file=is_valid_file,
        )

        self.redefine_class_to_idx()

        self.transform_ = transform
        self.target_transform_ = target_transform
        self.train = train
        self.return_path = return_path

        if zeroshot_split:
            indices_to_use = None   # Use all examples
            unique_classes = set(self.targets)
            # choose 25% of the classes for testing (with seed 77)
            random.seed(77)
            test_classes = set(random.sample(unique_classes, int(len(unique_classes) * 0.25)))
            train_classes = set(unique_classes) - test_classes
            self.target_classes = train_classes if self.train else test_classes

            # reindex the classes to 0, 1, 2, ...
            self.zs_class2class_id = dict(enumerate(self.target_classes))
            self.class_id2zs_class = {v: k for k, v in self.zs_class2class_id.items()}

            self.imgs = self.samples = [(img, self.class_id2zs_class[id]) for img, id in self.imgs if id in self.target_classes]
            self.img_paths = [img_path for img_path, _ in self.imgs]
            _, self.targets = list(zip(*self.samples))
        else:
            self.target_classes = set(self.targets)
            self.zs_class2class_id = None
            self.class_id2zs_class = None

            # obtain sample ids filtered by split
            path_to_splits = os.path.join(root, 'train_test_split.txt')
            indices_to_use = []
            with open(path_to_splits, 'r') as in_file:
                for line in in_file:
                    idx, use_train = line.strip('\n').split(' ', 2)
                    if bool(int(use_train)) == self.train:
                        indices_to_use.append(int(idx))

            # obtain filenames of images
            path_to_index = os.path.join(root, 'images.txt')
            filenames_to_use = []
            with open(path_to_index, 'r') as in_file:
                for line in in_file:
                    idx, fn = line.strip('\n').split(' ', 2)
                    if int(idx) in indices_to_use:
                        filenames_to_use.append(fn)

            img_paths_cut = {'/'.join(img_path.rsplit('/', 2)[-2:]): idx for idx, (img_path, lb) in enumerate(self.imgs)}
            imgs_to_use = [self.imgs[img_paths_cut[fn]] for fn in filenames_to_use]

            _, targets_to_use = list(zip(*imgs_to_use))

            self.imgs = self.samples = imgs_to_use
            self.img_paths = [img_path for img_path, _ in self.imgs]
            self.targets = targets_to_use

        # Data augmentation
        self.tri_aug = TrivialAugmentation(operations=[
            # 'ShearX',       # view point rotation
            # 'ShearY',
            # 'TranslateX',   # shifting
            # 'TranslateY',
            # 'Rotate',       # image rotation
            'Brightness',
            'Color',
            'Contrast',
            'Sharpness',
            # 'Posterize',  # quantize the color in 2^n, does not support float image
            # 'Solarize',   # invert the color below a threshold
            'AutoContrast',
            # 'Equalize'    # histogram equalization (the color will change), does not support float image
        ])

        self._load_meta(indices_to_use)

    def _load_meta(self, indices_to_use):
        # load part names idx to dict
        with open(os.path.join(self.root_dir, 'parts', 'parts.txt'), 'r') as f:
            self.dataset_parts = {int(line.split(' ')[0]): " ".join(line.split(' ')[1:]) for line in f.read().splitlines()}

        # load train test split
        with open(os.path.join(self.root_dir, 'train_test_split.txt'), 'r') as f:
            self.train_test_split = {int(line.split(' ')[0]): bool(int(line.split(' ')[1])) for line in f.readlines()}

        # load images.txt to get image id and corresponding image path
        with open(os.path.join(self.root_dir, 'images.txt'), 'r') as f:
            self.image_paths = {int(line.split(' ')[0]): line.split(' ')[1].strip('\n') for line in f.readlines()}

        # load part locations, e.g., {img_id: [[x, y, visible],  ...]], ...} note that each image has 15 parts
        self.part_locs = {}
        with open(os.path.join(self.root_dir, 'parts', 'part_locs.txt'), 'r') as f:
            for line in f.read().splitlines():
                img_id, part_id, x, y, visible = line.split(' ')
                self.part_locs.setdefault(int(img_id), []).append([int(float(x)), int(float(y)), int(float(visible))])

        # load bounding boxes
        with open(os.path.join(self.root_dir, 'bounding_boxes.txt'), 'r') as f:
            self.bounding_boxes = {int(line.split(' ')[0]): [int(float(x)) for x in line.split(' ')[1:]] for line in f.readlines()}

        # replace the img_id with image name (no extension)
        self.sample_parts = {}
        self.image_boxes = {}
        for img_id, image_path in self.image_paths.items():
            image_name = os.path.splitext(os.path.basename(image_path))[0]
            self.sample_parts[image_name] = self.part_locs[img_id]
            self.image_boxes[image_name] = self.bounding_boxes[img_id]

    def __getitem__(self, index):
        # generate one sample
        sample, target = super(CUB, self).__getitem__(index)

        if self.train:
            sample = self.tri_aug(sample)

        if self.transform_ is not None:
            sample = self.transform_(images=sample, return_tensor='pt')
        if self.target_transform_ is not None:
            target = self.target_transform_(images=target, return_tensor='pt')

        if self.return_path:
            return sample, target, self.img_paths[index]

        return sample, target

    def redefine_class_to_idx(self):
        adjusted_dict = {}
        for k, v in self.class_to_idx.items():
            k = k.split('.')[-1].replace('_', ' ')
            split_key = k.split(' ')
            if len(split_key) > 2:
                k = '-'.join(split_key[:-1]) + " " + split_key[-1]
            adjusted_dict[k] = v
        self.class_to_idx = adjusted_dict


class CUBBoxed(datasets.ImageFolder):
    """
    Wrapper for the CUB-200-2011 dataset.
    Method DatasetBirds.__getitem__() returns tuple of image and its corresponding label.
    Dataset per https://github.com/slipnitskaya/caltech-birds-advanced-classification
    """
    allowed_keys = ['crop', 'box_dir', 'return_path', 'trivial_aug', 'ops', 'high_res', 'n_pixel', 'return_mask']
    def __init__(self,
                 root: str,
                 transform: transforms.Compose= None,
                 target_transform: transforms.Compose= None,
                 loader= datasets.folder.default_loader,
                 is_valid_file= None,
                 train: bool = True,
                 **kwargs):

        img_root = os.path.join(root, 'images')
        # img_root = os.path.join(root, 'images_synthetic')
        self.__dict__.update((k, v) for k, v in kwargs.items() if k in self.allowed_keys)
        if hasattr(self, 'high_res') and self.high_res:
            img_root = os.path.join(root, 'images_upsampler')

        super(CUBBoxed, self).__init__(
            root=img_root,
            transform=None,
            target_transform=None,
            loader=loader,
            is_valid_file=is_valid_file,
        )

        self.redefine_class_to_idx()

        self.transform_ = transform
        self.target_transform_ = target_transform
        self.train = train
        self.totensor = transforms.ToTensor()

        self.resize = transforms.Resize((self.n_pixel, self.n_pixel), interpolation=transforms.InterpolationMode.BICUBIC, antialias=True, max_size=None) # change this may result in better performance

        # obtain sample ids filtered by split
        path_to_splits = os.path.join(root, 'train_test_split.txt')
        indices_to_use = []
        with open(path_to_splits, 'r') as in_file:
            for line in in_file:
                idx, use_train = line.strip('\n').split(' ', 2)
                if bool(int(use_train)) == self.train:
                    indices_to_use.append(int(idx))

        # obtain filenames of images
        path_to_index = os.path.join(root, 'images.txt')
        filenames_to_use = set()
        with open(path_to_index, 'r') as in_file:
            for line in in_file:
                idx, fn = line.strip('\n').split(' ', 2)
                if int(idx) in indices_to_use:
                    filenames_to_use.add(fn)

        img_paths_cut = {'/'.join(img_path.rsplit('/', 2)[-2:]): idx for idx, (img_path, lb) in enumerate(self.imgs)}
        imgs_to_use = [self.imgs[img_paths_cut[fn]] for fn in filenames_to_use]

        _, targets_to_use = list(zip(*imgs_to_use))

        self.imgs = self.samples = imgs_to_use
        self.targets = targets_to_use
        
        mean = list(self.transform_.transforms)[-1].mean
        self.gray_image = torch.ones((1, 3, self.n_pixel, self.n_pixel), dtype=torch.float32) * torch.tensor(mean)[None, :, None, None]

    def __getitem__(self, index):
        # generate one sample
        path, target = self.samples[index]
        sample = self.loader(path)
        sample_size = torch.tensor(sample.size[::-1]) # (w, h) -> (h, w)
        hight, width = sample_size
        image_id = os.path.splitext(os.path.basename(path))[0]
        img_tensor = self.totensor(sample)
        # repeat samples
        boxes = torch.load(os.path.join(self.box_dir, image_id + '.pth'))["boxes_info"]
        sample = torch.cat([self.resize(img_tensor).unsqueeze(0)] * len(boxes))
        box_labels = torch.ones(len(boxes))
        masks = torch.zeros((len(boxes), self.n_pixel, self.n_pixel), dtype=torch.uint8)
        
        if hasattr(self, 'crop') or hasattr(self, 'return_mask'):
            crop_samples = []
            for idx, (part_name, box) in enumerate(boxes.items()):
                # boxes = {'Part name0': [x0, y0, x1, y1], 'Part name1': [x0, y0, x1, y1], ...}
                if self.crop:
                    if sum(box) == 0:
                        box_image = self.gray_image
                    else:
                        box_image = self.resize(img_tensor[:, box[1]:box[3], box[0]:box[2]].unsqueeze(0))
                    crop_samples.append(box_image)
                    
                if self.return_mask:
                    # scale box from image_size (h, w) to (n_pixel, n_pixel)
                    scale_factor = torch.tensor([self.n_pixel / width, self.n_pixel / hight, self.n_pixel / width, self.n_pixel / hight])
                    box = (torch.tensor(box) * scale_factor).int()
                    masks[idx, box[1]:box[3], box[0]:box[2]] = 1
                
                if sum(box) == 0:
                    box_labels[idx] = 0
            crop_sample = torch.cat(crop_samples)
        if len(crop_samples) == 0:
            crop_samples = torch.zeros_like(sample)


        if self.transform_ is not None:
            sample = self.transform_(sample)
            if hasattr(self, 'crop') and self.crop:
                crop_sample = self.transform_(crop_sample)

        # return (sample, target, path, sample_size) if hasattr(self, 'return_path') and self.return_path else (sample, target)
        if hasattr(self, 'return_path') and self.return_path:
            return sample, target, path, sample_size
        # if hasattr(self, 'return_mask') and self.return_mask:
        #     return sample, target, masks
        return sample, target, masks, crop_sample, path, box_labels

    def redefine_class_to_idx(self):
        adjusted_dict = {}
        for k, v in self.class_to_idx.items():
            k = k.split('.')[-1].replace('_', ' ')
            split_key = k.split(' ')
            if len(split_key) > 2:
                k = '-'.join(split_key[:-1]) + " " + split_key[-1]
            adjusted_dict[k] = v
        self.class2idx = adjusted_dict
        self.idx2class = {v: k for k, v in self.class_to_idx.items()}
src/data_loader/cub.py
anguyen8/peeb
import random

from torchvision.datasets.folder import default_loader
from torchvision.datasets import ImageFolder, VisionDataset, ImageNet, Places365, EuroSAT, DTD, Food101, OxfordIIITPet
from imagenetv2_pytorch import ImageNetV2Dataset as ImageNetV2
from torch.utils.data import Subset

from .cub import CUB
from .nabirds import NABirdsDataset
from .bird_soup import BirdSoup
from configs import *


class DatasetWrapper(VisionDataset):
    allowed_keywords = ["idx2name", "name2idx", "target_classes", "image_type_cub"]

    def __init__(self, dataset_name, transform, distortion=None, samples_per_class: int = None, random_seed: int = 42, **kwargs):
        self.transform = transform
        self.dataset_name = dataset_name
        self.__dict__.update((k, v) for k, v in kwargs.items() if k in self.allowed_keywords)

        # Used for ImageNet-C only
        results = distortion.split(":") if distortion is not None else [None, None]
        self.distortion_name, self.distortion_severity = results[0], results[1]

        self.dataset, self.n_classes = self.load_dataset(dataset_name, samples_per_class, random_seed)

    def load_dataset(self, dataset_name, samples_per_class, random_seed: int = 42):
        random.seed(random_seed)

        if dataset_name == DS_IMAGENET:
            dataset = ImageNet(IMAGENET_DIR, split='val', transform=self.transform)
        elif dataset_name == DS_IMAGENET_V2:
            dataset = ImageNetV2(location=IMAGENET_V2_DIR, transform=self.transform)
        elif dataset_name == DS_IMAGENET_A:
            dataset = ImageFolder(root=IMAGENET_A_DIR, transform=self.transform)
        elif dataset_name == DS_IMAGENET_C:
            dataset = ImageFolder(root=f"{IMAGENET_C_DIR}/{self.distortion_name}/{self.distortion_severity}/", transform=self.transform)
        elif dataset_name == DS_PLACES365:
            dataset = Places365(PLACES365_DIR, split='val', transform=self.transform)
        elif dataset_name == DS_CUB:
            dataset = CUB(CUB_DIR, train=False, transform=self.transform, image_type=self.image_type_cub)
        elif dataset_name == DS_NABIRDS:
            dataset = NABirdsDataset(NABIRDS_DIR, transform=self.transform, train=False)
        elif dataset_name == DS_BIRD_SOUP:
            meta_path = BIRD_SOUP_META_PATH_ALL   # BIRD_SOUP_META_PATH_ALL / BIRD_SOUP_META_PATH_TRAIN / BIRD_SOUP_META_PATH_TEST
            dataset = BirdSoup(BIRD_SOUP_DIR, transform=self.transform, train=False, return_path=False, meta_path=meta_path)

        n_classes = len(dataset.classes) if dataset_name != DS_IMAGENET_V2 else 1000

        # Using Subset
        if samples_per_class is not None and samples_per_class > 0:
            class_counts = self.get_class_count(dataset_name, dataset)

            selected_indices = []
            total_samples_processed = 0
            for label, max_count in class_counts.items():
                selected_indices.extend(sorted(list(random.sample(list(range(total_samples_processed, total_samples_processed + max_count)), k=min(samples_per_class, max_count)))))
                total_samples_processed += max_count
            dataset = Subset(dataset, selected_indices)

        # Select subset if target_classes is specified and less than the total number of classes
        if hasattr(self, 'target_classes') and len(self.target_classes) < len(set(dataset.targets)):
            select_indices = [i for i, t in enumerate(dataset.targets) if t in self.target_classes]
            dataset = Subset(dataset, select_indices)

        return dataset, n_classes

    def __getitem__(self, idx):
        sample, target = self.dataset.__getitem__(idx)
        path, image_id = "", ""

        if self.dataset_name in [DS_IMAGENET, DS_IMAGENET_A, DS_IMAGENET_C, DS_PLACES365, DS_CUB, DS_NABIRDS, DS_BIRD_SOUP]:
            path, label = self.dataset.dataset.imgs[self.dataset.indices[idx]] if isinstance(self.dataset, Subset) else self.dataset.imgs[idx]
            image_id = ".".join(path.split("/")[-1].split(".")[:-1])
        elif self.dataset_name == DS_IMAGENET_V2:
            posix_path = self.dataset.dataset.fnames[self.dataset.indices[idx]] if isinstance(self.dataset, Subset) else self.dataset.fnames[idx]
            path, image_id = str(posix_path), posix_path.stem

        image = default_loader(path)
        image_size = torch.tensor(image.size[::-1])  # (w, h) -> (1, h, w)

        return sample, target, path, image_id, image_size

    def __len__(self):
        return len(self.dataset)

    def get_class_count(self, dataset_name, dataset):
        class_counts = {}

        if dataset_name in [DS_IMAGENET, DS_IMAGENET_A, DS_IMAGENET_C, DS_PLACES365, DS_CUB]:
            for path, label in dataset.imgs:
                if label not in class_counts:
                    class_counts[label] = 0
                class_counts[label] += 1

        elif dataset_name == DS_IMAGENET_V2:
            for posix_path in dataset.fnames:
                label = posix_path.parts[-2]
                if label not in class_counts:
                    class_counts[label] = 0
                class_counts[label] += 1

        return class_counts

src/data_loader/dataset_wrapper.py
anguyen8/peeb
import os
import json
import torch
from PIL import Image

from torchvision import transforms
from torchvision.transforms import Compose
from torch.utils.data import Dataset
from torchvision.datasets.folder import default_loader

from src.data_loader.augmentation import TrivialAugmentation
from src.data_loader.nabirds_horn import load_bounding_box_annotations, load_part_annotations, load_part_names, load_class_names, \
     load_image_labels, load_image_paths, load_image_sizes, load_hierarchy, load_photographers, load_train_test_split

import random


class NABirdsDataset(Dataset):
    allowed_keys = ['crop', 'box_dir', 'return_path', 'trivial_aug', 'ops', 'high_res', 'n_pixel', 'return_mask', 'all_data', 'zeroshot_split']
    part_mapping = {'beak': [0], 'crown': [1], 'nape': [3], 'eyes': [3, 4], 'belly': [5], 'breast': [6], 'back': [7], 'tail': [8], 'wings': [9, 10]}

    def __init__(self, root_dir: str, transform: Compose = None, train: bool = True, zeroshot_split: bool = False, **kwargs):

        self.root_dir = root_dir
        self.zeroshot_split = zeroshot_split

        self.transform = transform
        self.train = train
        self.__dict__.update({k: v for k, v in kwargs.items() if k in self.allowed_keys})
        self.__dict__.update((k, v) for k, v in kwargs.items() if k in self.allowed_keys)

        self._load_meta()

        self.loader = default_loader
        self.totensor = transforms.ToTensor()

        # Data augmentation
        self.tri_aug = TrivialAugmentation(operations=[
            # 'ShearX',  # view point rotation
            # 'ShearY',
            # 'TranslateX',  # shifting
            # 'TranslateY',
            # 'Rotate',  # image rotation
            'Brightness',
            'Color',
            'Contrast',
            'Sharpness',
            # 'Posterize',  # quantize the color in 2^n, does not support float image
            # 'Solarize',   # invert the color below a threshold
            'AutoContrast',
            # 'Equalize'    # histogram equalization (the color will change), does not support float image
        ])

    def _load_meta(self):

        # ================== Load meta data ==================
        # Loading code adopt form NABirds sample script with minor modification
        self.image_paths = load_image_paths(self.root_dir, path_prefix='images')
        self.image_sizes = load_image_sizes(self.root_dir)
        self.image_bboxes = load_bounding_box_annotations(self.root_dir)
        self.image_parts = load_part_annotations(self.root_dir)
        self.image_class_labels = load_image_labels(self.root_dir)

        # Load in the class data
        self.class_names = load_class_names(self.root_dir)
        self.class_hierarchy = load_hierarchy(self.root_dir)

        # Load in the part data
        self.dataset_parts = load_part_names(self.root_dir)
        self.part_ids = list(self.dataset_parts.keys())
        self.part_ids = sorted(self.part_ids)

        # Load in the train / test split
        self.train_images, self.test_images = load_train_test_split(self.root_dir)

        # ===================================================
        # replace part names 'bill' with 'beak'
        # (this is the parts from the dataset, not to confused with part_name used for training, e.g., self.part_names)
        self.dataset_parts = {0: 'beak', 1: 'crown', 2: 'nape', 3: 'left eye', 4: 'right eye', 5: 'belly',
                              6: 'breast', 7: 'back', 8: 'tail', 9: 'left wing', 10: 'right wing'}

        all_samples = self.train_images + self.test_images
        all_classes = [int(self.image_class_labels[image_id]) for image_id in all_samples]
        unique_classes = set(all_classes)
        unique_class_names = [self.class_names[str(class_id)] for class_id in unique_classes]

        # there are 555 classes in total, re-index them from 0 to 554
        self.class2idx = {v: k for k, v in enumerate(unique_class_names)}
        self.idx2class = dict(enumerate(unique_class_names))
        self.nabirds_idx2class_idx = dict(zip(unique_classes, self.idx2class))
        self.class_idx2nabirds_idx = dict(zip(self.idx2class, unique_classes))

        # training samples
        self.target_classes = unique_classes
        self.zs_class2class_id = None
        self.class_id2zs_class = None

        if self.zeroshot_split:
            # choose 25% of the classes for testing (with seed 77)
            random.seed(77)     # Fixed random seed
            test_classes = set(random.sample(unique_classes, int(len(unique_classes) * 0.25)))
            train_classes = set(unique_classes) - test_classes
            all_imgs = self.train_images + self.test_images

            self.target_classes = train_classes if self.train else test_classes
            # reindex the classes to 0, 1, 2, ...
            self.zs_class2class_id = dict(enumerate(self.target_classes))
            self.class_id2zs_class = {v: k for k, v in self.zs_class2class_id.items()}
            image_ids = [image_id for image_id in all_imgs if int(self.image_class_labels[image_id]) in self.target_classes]
        elif hasattr(self, 'all_data') and self.all_data:
            image_ids = all_samples
        else:
            image_ids = self.train_images if self.train else self.test_images

        self.samples = []
        self.targets = []
        self.sample_parts = {}
        self.image_boxes = {}

        for image_id in image_ids:
            image_path = self.image_paths[image_id]
            image_name = os.path.splitext(os.path.basename(image_path))[0]
            class_label = self.nabirds_idx2class_idx[int(self.image_class_labels[image_id])]

            if self.class_id2zs_class is not None:
                class_label = self.class_id2zs_class[class_label]

            self.targets.append(class_label)
            self.samples.append((os.path.join(self.root_dir, image_path), class_label))
            self.sample_parts[image_name] = self.image_parts[image_id]
            self.image_boxes[image_name] = self.image_bboxes[image_id]

        self.classes = list(self.class2idx.keys())
        self.imgs = self.samples

    def __len__(self):
        return len(self.samples)

    # overwrite __getitem__
    def __getitem__(self, idx):
        image_path, target = self.samples[idx]

        sample = self.loader(image_path)
        sample_size = torch.tensor(sample.size[::-1])  # (h, w)

        if self.tri_aug is not None and self.train:
            sample = self.tri_aug(sample)

        if self.transform is not None:
            sample = self.transform(images=sample, return_tensor='pt')

        if hasattr(self, 'return_path') and self.return_path:
            return sample, target, image_path

        return sample, target
src/data_loader/nabirds.py
anguyen8/peeb
"""_summary_
The sample script from NABirds dataset. Source: https://dl.allaboutbirds.org/nabirds. Modified to adapt python3.
"""

import os
import sys

from matplotlib import pyplot as plt
import random
import math

try:
    import networkx as nx

    has_networkx = True
except:
    has_networkx = False

import os


def load_bounding_box_annotations(dataset_path=''):
    bboxes = {}

    with open(os.path.join(dataset_path, 'bounding_boxes.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            image_id = pieces[0]
            bbox = list(map(int, pieces[1:]))
            bboxes[image_id] = bbox

    return bboxes


def load_part_annotations(dataset_path=''):
    parts = {}

    with open(os.path.join(dataset_path, 'parts/part_locs.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            image_id = pieces[0]
            parts.setdefault(image_id, [0] * 11)
            part_id = int(pieces[1])
            parts[image_id][part_id] = list(map(int, pieces[2:]))

    return parts


def load_part_names(dataset_path=''):
    names = {}

    with open(os.path.join(dataset_path, 'parts/parts.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            part_id = int(pieces[0])
            names[part_id] = ' '.join(pieces[1:])

    return names


def load_class_names(dataset_path=''):
    names = {}

    with open(os.path.join(dataset_path, 'classes.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            class_id = pieces[0]
            names[class_id] = ' '.join(pieces[1:])

    return names


def load_image_labels(dataset_path=''):
    labels = {}

    with open(os.path.join(dataset_path, 'image_class_labels.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            image_id = pieces[0]
            class_id = pieces[1]
            labels[image_id] = class_id

    return labels


def load_image_paths(dataset_path='', path_prefix=''):
    paths = {}

    with open(os.path.join(dataset_path, 'images.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            image_id = pieces[0]
            path = os.path.join(path_prefix, pieces[1])
            paths[image_id] = path

    return paths


def load_image_sizes(dataset_path=''):
    sizes = {}

    with open(os.path.join(dataset_path, 'sizes.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            image_id = pieces[0]
            width, height = map(int, pieces[1:])
            sizes[image_id] = [width, height]

    return sizes


def load_hierarchy(dataset_path=''):
    parents = {}

    with open(os.path.join(dataset_path, 'hierarchy.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            child_id, parent_id = pieces
            parents[child_id] = parent_id

    return parents


def load_photographers(dataset_path=''):
    photographers = {}
    with open(os.path.join(dataset_path, 'photographers.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            image_id = pieces[0]
            photographers[image_id] = ' '.join(pieces[1:])

    return photographers


def load_train_test_split(dataset_path=''):
    train_images = []
    test_images = []

    with open(os.path.join(dataset_path, 'train_test_split.txt')) as f:
        for line in f:
            pieces = line.strip().split()
            image_id = pieces[0]
            is_train = int(pieces[1])
            if is_train:
                train_images.append(image_id)
            else:
                test_images.append(image_id)
    return train_images, test_images


if __name__ == '__main__':

    if len(sys.argv) > 1:
        dataset_path = sys.argv[1]
    else:
        dataset_path = ''

    if len(sys.argv) > 2:
        image_path = sys.argv[2]
    else:
        image_path = 'images'

    # Load in the image data
    # Assumes that the images have been extracted into a directory called "images"
    image_paths = load_image_paths(dataset_path, path_prefix=image_path)
    image_sizes = load_image_sizes(dataset_path)
    image_bboxes = load_bounding_box_annotations(dataset_path)
    image_parts = load_part_annotations(dataset_path)
    image_class_labels = load_image_labels(dataset_path)

    # Load in the class data
    class_names = load_class_names(dataset_path)
    class_hierarchy = load_hierarchy(dataset_path)

    # Load in the part data
    part_names = load_part_names(dataset_path)
    part_ids = part_names.keys()
    part_ids = sorted(part_ids)

    # Load in the photographers
    photographers = load_photographers(dataset_path)

    # Load in the train / test split
    train_images, test_images = load_train_test_split(dataset_path)

    # Visualize the images and their annotations
    image_identifiers = list(image_paths.keys())
    random.shuffle(image_identifiers)
    for image_id in image_identifiers:

        image_path = image_paths[image_id]
        image = plt.imread(image_path)
        bbox = image_bboxes[image_id]
        parts = [list(part) for part in image_parts[image_id]]
        class_label = image_class_labels[image_id]
        class_name = class_names[class_label]

        # Render the image
        plt.close(1)
        fig = plt.figure(1, figsize=(16, 12))
        plt.imshow(image)
        plt.title("Image ID: %s \n Class Label: %s  [ID: %s]" % (image_id, class_name, class_label))

        # Render the bounding box annotations
        bbox_x, bbox_y, bbox_width, bbox_height = bbox
        currentAxis = plt.gca()
        currentAxis.add_patch(plt.Rectangle((bbox_x, bbox_y), bbox_width, bbox_height, facecolor="green", alpha=0.3))

        # Render the part annotations
        if has_networkx:
            # Use networkx spring layout

            G = nx.Graph()
            part_data = []
            initial_positions = {}
            for part_id, part in enumerate(parts):
                x, y, v = part
                if v:
                    part_str = 'part_%d' % (part_id,)
                    label_str = 'label_%d' % (part_id,)

                    G.add_node(part_str)
                    G.add_node(label_str)
                    G.add_edge(part_str, label_str)

                    part_data.append(part_str)

                initial_positions[part_str] = (x, y)
                initial_positions[label_str] = (x, y)

        positions = nx.spring_layout(G, dim=2, k=20.0, pos=initial_positions, fixed=part_data)

        for part_id, part in enumerate(parts):
            x, y, v = part
            if v:
                plt.plot(x, y, 'o')

                label_str = 'label_%d' % (part_id,)
                label_position = positions[label_str]
                label_xy = (label_position[0] * bbox_width + bbox_x, label_position[1] * bbox_height + bbox_y)
                plt.annotate(
                    part_names[part_id],
                    xy=(x, y),
                    xycoords='data',
                    xytext=label_xy,
                    textcoords='data',
                    ha='right',
                    va='bottom',
                    bbox=dict(boxstyle='round,pad=0.5',
                              fc='yellow',
                              alpha=0.5),
                    arrowprops=dict(arrowstyle='->',
                                    connectionstyle='arc3,rad=0')
                )

        else:
            # just a basic label annotation for the part locations
            offset = (2 * math.pi) / len(parts)
            for part_id, part in enumerate(parts):
                x, y, v = part
                if v:
                    # try to offset the part labels so that they don't overlap too badly
                    dist = random.randint(35, 65)
                    offset_x = dist * math.cos(offset * part_id)
                    offset_y = dist * math.sin(offset * part_id)

                plt.plot(x, y, 'o')
                plt.annotate(
                    part_names[part_id],
                    xy=(x, y),
                    xytext=(offset_x, offset_y),
                    textcoords='offset points',
                    ha='right',
                    va='bottom',
                    bbox=dict(boxstyle='round,pad=0.5',
                              fc='yellow',
                              alpha=0.5),
                    arrowprops=dict(arrowstyle='->',
                                    connectionstyle='arc3,rad=0')
                )

        plt.show(block=False)
        # Print some info about the image and the annotations:
        print(f"Image ID: {image_id}")
        print(f"Image Path: {image_path}")
        print(f"Image Dimensions: {image_sizes[image_id][0]} x {image_sizes[image_id][1]}")
        print(f"Photographer: {photographers[image_id]}")
        print(f"Split: {'train' if image_id in train_images else 'test'}")
        print("Class Label Path:")

        label_path = []
        current_class_label = class_label
        while current_class_label in class_hierarchy:
            label_path.append((class_names[current_class_label], current_class_label))
            current_class_label = class_hierarchy[current_class_label]
        label_path.append((class_names[current_class_label], current_class_label))  # add the root node
        label_path.reverse()
        for depth, (class_name, class_id) in enumerate(label_path):
            line = "\t" * (depth + 1)
            print(f"{line}|- {class_name} \t[ID: {class_id}]")

        print(f"Bounding Box Dimensions: {bbox_width} x {bbox_height}")
        print(f"{'Part Name':<10}\t{'Visibility':<1}")
        for part_id in part_ids:
            x, y, v = parts[part_id]
            print(f"{part_names[part_id]:<10}\t{'Y' if v else 'N'}")

        # Keep showing images until a button other than Enter is pressed
        ri = input("Press Enter to continue. Or press any key and then enter to quit...")
        if ri != '':
            break
        print()
        print()

src/data_loader/nabirds_horn.py
anguyen8/peeb
import copy
from typing import Optional, Tuple, Union, Any

import torch
import torch.nn.functional as F
import numpy as np
from torch import nn
from transformers import OwlViTConfig
from transformers.models.owlvit.modeling_owlvit import OwlViTVisionTransformer
from transformers.models.detr.modeling_detr import generalized_box_iou, center_to_corners_format
from transformers.image_transforms import corners_to_center_format

from scipy.optimize import linear_sum_assignment
from transformers.utils import requires_backends
from typing import Dict, List, Optional, Tuple
from torch import Tensor, nn


class OwlViTClassPredictionHead(nn.Module):
    def __init__(self, config: OwlViTConfig):
        super().__init__()

        out_dim = config.text_config.hidden_size
        self.query_dim = config.vision_config.hidden_size

        self.dense0 = nn.Linear(self.query_dim, out_dim)
        self.logit_shift = nn.Linear(self.query_dim, 1)
        self.logit_scale = nn.Linear(self.query_dim, 1)
        self.elu = nn.ELU()

    def forward(
        self,
        image_embeds: torch.FloatTensor,
        query_embeds: Optional[torch.FloatTensor],
        query_mask: Optional[torch.Tensor],
    ) -> Tuple[torch.FloatTensor]:
        image_class_embeds = self.dense0(image_embeds)
        if query_embeds is None:
            device = image_class_embeds.device
            batch_size, num_patches = image_class_embeds.shape[:2]
            pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device)
            return (pred_logits, image_class_embeds)

        # Normalize image and text features
        image_class_embeds = F.normalize(image_class_embeds, dim=-1) + 1e-6
        query_embeds = F.normalize(query_embeds, dim=-1) + 1e-6

        # Get class predictions
        pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)

        # Apply a learnable shift and scale to logits
        logit_shift = self.logit_shift(image_embeds)
        logit_scale = self.logit_scale(image_embeds)
        logit_scale = self.elu(logit_scale) + 1
        pred_logits = (pred_logits + logit_shift) * logit_scale

        if query_mask is not None:
            if query_mask.ndim > 1:
                query_mask = torch.unsqueeze(query_mask, dim=-2)

            pred_logits = pred_logits.to(torch.float64)
            pred_logits = torch.where(query_mask == 0, -1e6, pred_logits)
            pred_logits = pred_logits.to(torch.float32)

        return (pred_logits, image_class_embeds)


class OwlViTBoxPredictionHead(nn.Module):
    def __init__(self, config: OwlViTConfig, num_layers: int = 3):
        super().__init__()

        print("OwlViTBoxPredictionHead: Number of layers: ", num_layers)

        self.num_layers = num_layers
        width = config.vision_config.hidden_size
        self.dense0 = nn.Linear(width, width)
        self.dense1 = nn.Linear(width, width)
        self.dense2 = nn.Linear(width, width)
        self.dense3 = nn.Linear(width, width)
        self.dense4 = nn.Linear(width, width)
        self.dense5 = nn.Linear(width, width)
        self.gelu = nn.GELU()
        self.dense6 = nn.Linear(width, 4)

    def forward(self, image_features: torch.Tensor) -> torch.FloatTensor:
        output = self.dense0(image_features)
        output = self.gelu(output)
        output = self.dense1(output)
        output = self.gelu(output)

        if self.num_layers == 3:
            output = self.dense6(output)
            output = self.gelu(output)
            return output

        output = self.dense2(output)
        output = self.gelu(output)
        output = self.dense3(output)
        output = self.gelu(output)

        if self.num_layers == 5:
            output = self.dense6(output)
            output = self.gelu(output)
            return output

        output = self.dense4(output)
        output = self.gelu(output)
        output = self.dense5(output)
        output = self.gelu(output)
        output = self.dense6(output)
        output = self.gelu(output)

        return output


class OwlViTPredictionHead(nn.Module):
    def __init__(self, config: OwlViTConfig, num_classes: int):
        super().__init__()

        out_dim = config.text_config.hidden_size
        self.query_dim = config.vision_config.hidden_size

        self.num_classes = num_classes

        # Classification projection in paper
        self.mlp_image = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=self.query_dim, out_features=self.query_dim),
            nn.GELU(),
            nn.Linear(in_features=self.query_dim, out_features=self.query_dim),
            nn.GELU(),
            nn.Linear(in_features=self.query_dim, out_features=out_dim),
            nn.GELU(),
        )

    def forward(self,
                image_embeds: torch.FloatTensor,
                query_embeds: torch.FloatTensor,
                topk_idxs: torch.FloatTensor,
    ) -> Tuple[torch.FloatTensor]:

        # Get class predictions: topk_idxs (batch_size, n_parts, 1), one_hot (batch_size, n_parts, n_patches*n_patches)
        topk_idxs = torch.swapaxes(topk_idxs, 1, 2)
        one_hot = torch.zeros(topk_idxs.shape[0], topk_idxs.shape[1], image_embeds.shape[1]).to(image_embeds.device).scatter_(2, topk_idxs, 1)
        batch_size, n_parts = one_hot.shape[0], one_hot.shape[1]

        # (batch_size, n_parts, 3600, 1) * (batch_size, 1, 3600, 1024) = (batch_size, n_parts, 3600, 1024).sum(dim=-2)
        image_embeds = (one_hot.unsqueeze(-1) * image_embeds.unsqueeze(1)).sum(dim=-2)

        # image_embeds = self.dense0(image_embeds)            # (batch_size, n_patches, 1024) --> (.., .., 768)
        image_embeds = self.mlp_image(image_embeds.view(-1, image_embeds.shape[-1])).view(batch_size, n_parts, -1)
        query_embeds = query_embeds.view(batch_size, -1, query_embeds.shape[-1])

        # Normalize image and text features
        image_embeds = F.normalize(image_embeds, dim=-1)  # (batch_size, n_parts, 768)
        query_embeds = F.normalize(query_embeds, dim=-1)  # (batch_size, num_classes * n_parts, 768)

        # Send query_embeds to the same device as image_embeds
        query_embeds = query_embeds.to(image_embeds.device)

        # Shape: torch.Size([bs, num_boxes, num_classes * num_parts])
        image_text_logits = torch.einsum('bnd, bid -> bni', image_embeds, query_embeds)
        image_text_logits_reshaped = image_text_logits.view(-1, image_text_logits.shape[-1])

        # Shape: (bs, num_classes * num_parts, num_boxes) --> (bs, num_classes, num_parts, num_boxes)
        pred_logits = image_text_logits.swapaxes(axis0=1, axis1=2).view(batch_size, self.num_classes, n_parts, -1)
        pred_logits = torch.diagonal(pred_logits, dim1=-2, dim2=-1)     # --> torch.Size([bs, num_classes, 12])
        pred_logits = torch.sum(pred_logits, dim=-1)

        return (image_text_logits_reshaped, pred_logits)


class OwlViTForClassification(nn.Module):
    config_class = OwlViTConfig

    def __init__(self, owlvit_det_model, num_classes, num_parts, weight_dict, device, freeze_box_heads=False,
                 train_box_heads_only=False, network_type=None, classification_loss=None, logits_from_teacher=False,
                 finetuning=None, alpha=0.25, gamma=2.0):
        super(OwlViTForClassification, self).__init__()

        self.config = owlvit_det_model.config
        self.num_classes = num_classes
        self.num_parts = num_parts
        self.device = device

        self.alpha = alpha
        self.gamma = gamma
        self.sigmoid = nn.Sigmoid()
        self.ce_loss = torch.nn.CrossEntropyLoss()
        self.focal_loss = FocalLoss(alpha=self.alpha, gamma=self.gamma, reduction='mean')

        # Use CE loss for classification OR only train with contrastive loss
        self.network_type = network_type
        self.classification_loss = classification_loss
        self.logits_from_teacher = logits_from_teacher
        self.finetuning = finetuning

        # Initialize OwlViT model from the teacher model
        self.owlvit = copy.deepcopy(owlvit_det_model.owlvit)
        self.layer_norm = copy.deepcopy(owlvit_det_model.layer_norm)

        # For image-level classification
        self.cls_head = OwlViTPredictionHead(self.config, self.num_classes)

        # For box prediction
        self.box_head = copy.deepcopy(owlvit_det_model.box_head)

        # For box-level classification
        self.class_head = OwlViTClassPredictionHead(self.config)
        self.class_head.dense0.load_state_dict(owlvit_det_model.class_head.dense0.state_dict())
        self.class_head.logit_shift.load_state_dict(owlvit_det_model.class_head.logit_shift.state_dict())
        self.class_head.logit_scale.load_state_dict(owlvit_det_model.class_head.logit_scale.state_dict())

        # OwlViT: set equal weights for the bounding box, gIoU and classification losses
        # self.matcher = DetrHungarianMatcher(class_cost=1, bbox_cost=1, giou_cost=1)

        # Losses for the criterion in DETR/OwlViT
        self.weight_dict = weight_dict
        losses = ["cardinality"]
        losses += ["boxes"] if weight_dict["loss_bbox"] > 0 else []
        losses += ["labels"] if weight_dict["loss_ce"] > 0 else []

        self.criterion = DetrLoss(
            matcher=None,
            num_parts=self.num_parts,
            eos_coef=0.1,   # Following facebook/detr-resnet-50
            losses=losses,
        )

        self.freeze_parameters(freeze_box_heads, train_box_heads_only, finetuning)
        del owlvit_det_model

    def init_weights_he(self, m):
        if type(m) == nn.Linear:
            nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
            m.bias.data.fill_(0.01)

    def update_box_head(self, num_layers):
        # TODO: FORCE INITIALIZATION OF BOX HEAD's WEIGHTS
        self.box_head = OwlViTBoxPredictionHead(self.config, num_layers=num_layers)
        self.box_head.apply(self.init_weights_he)

    def freeze_parameters(self, freeze_box_heads, train_box_heads_only, finetuning):
        # OwlViT's text encoder is frozen by default
        for param in self.owlvit.text_model.parameters():
            param.requires_grad = False
        for param in self.owlvit.text_projection.parameters():
            param.requires_grad = False

        # SKIP finetuning box heads
        if freeze_box_heads:
            for param in self.box_head.parameters():
                param.requires_grad = False
            for param in self.class_head.parameters():
                param.requires_grad = False

            # Freeze OwlViT vision encoder
            if finetuning == "mlp_only":
                for param in self.owlvit.parameters():
                    param.requires_grad = False
                for param in self.layer_norm.parameters():
                    param.requires_grad = False

        # SKIP finetuning vision encoder and MLP head for classification --> Adjust weights of box heads only
        if train_box_heads_only:
            for param in self.owlvit.parameters():
                param.requires_grad = False
            for param in self.layer_norm.parameters():
                param.requires_grad = False
            for param in self.cls_head.parameters():
                param.requires_grad = False

    def update_num_classes(self, num_classes):
        self.num_classes = num_classes
        self.cls_head.num_classes = num_classes

    def image_text_embedder(self,
        input_ids: torch.Tensor,
        pixel_values: torch.FloatTensor,
        attention_mask: torch.Tensor,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
    ) -> Tuple[torch.FloatTensor]:

        # Encode text and image
        outputs = self.owlvit(
            pixel_values=pixel_values,
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )

        # Get image embeddings
        last_hidden_state = outputs.vision_model_output[0]      # 0: last_hidden_state; 1: pooled_output
        image_embeds = self.owlvit.vision_model.post_layernorm(last_hidden_state)

        # Resize class token
        new_size = tuple(np.array(image_embeds.shape) - np.array((0, 1, 0)))
        class_token_out = torch.broadcast_to(image_embeds[:, :1, :], new_size)

        # Merge image embedding with class tokens
        image_embeds = image_embeds[:, 1:, :] * class_token_out
        image_embeds = self.layer_norm(image_embeds)

        # Resize to [batch_size, num_patches, num_patches, hidden_size]
        new_size = (
            image_embeds.shape[0],
            int(np.sqrt(image_embeds.shape[1])),
            int(np.sqrt(image_embeds.shape[1])),
            image_embeds.shape[-1],
        )
        image_embeds = image_embeds.reshape(new_size)
        text_embeds = outputs[-4]

        return (text_embeds, image_embeds, outputs)

    def image_embedder(
        self,
        pixel_values: torch.FloatTensor
    ) -> Tuple[torch.FloatTensor]:

        # Get OwlViTModel vision embeddings (same as CLIP)
        vision_outputs = self.owlvit.vision_model(pixel_values=pixel_values, return_dict=True)

        # Apply post_layernorm to last_hidden_state, return non-projected output
        last_hidden_state = vision_outputs[0]
        image_embeds = self.owlvit.vision_model.post_layernorm(last_hidden_state)

        # Resize class token
        new_size = tuple(np.array(image_embeds.shape) - np.array((0, 1, 0)))
        class_token_out = torch.broadcast_to(image_embeds[:, :1, :], new_size)

        # Merge image embedding with class tokens
        image_embeds = image_embeds[:, 1:, :] * class_token_out
        image_embeds = self.layer_norm(image_embeds)

        # Resize to [batch_size, num_patches, num_patches, hidden_size]
        new_size = (
            image_embeds.shape[0],
            int(np.sqrt(image_embeds.shape[1])),
            int(np.sqrt(image_embeds.shape[1])),
            image_embeds.shape[-1],
        )
        image_embeds = image_embeds.reshape(new_size)

        return (image_embeds, vision_outputs)

    def normalize_grid_corner_coordinates(self, feature_map: torch.FloatTensor):
        # Computes normalized xy corner coordinates from feature_map.
        if not feature_map.ndim == 4:
            raise ValueError("Expected input shape is [batch_size, num_patches, num_patches, hidden_dim]")

        device = feature_map.device
        num_patches = feature_map.shape[1]

        box_coordinates = np.stack(np.meshgrid(np.arange(1, num_patches + 1), np.arange(1, num_patches + 1)), axis=-1).astype(np.float32)
        box_coordinates /= np.array([num_patches, num_patches], np.float32)

        # Flatten (h, w, 2) -> (h*w, 2)
        box_coordinates = box_coordinates.reshape(box_coordinates.shape[0] * box_coordinates.shape[1], box_coordinates.shape[2])
        box_coordinates = torch.from_numpy(box_coordinates).to(device)

        return box_coordinates

    def compute_box_bias(self, feature_map: torch.FloatTensor) -> torch.FloatTensor:
        # The box center is biased to its position on the feature grid
        box_coordinates = self.normalize_grid_corner_coordinates(feature_map)
        box_coordinates = torch.clip(box_coordinates, 0.0, 1.0)

        # Unnormalize xy
        box_coord_bias = torch.log(box_coordinates + 1e-4) - torch.log1p(-box_coordinates + 1e-4)

        # The box size is biased to the patch size
        box_size = torch.full_like(box_coord_bias, 1.0 / feature_map.shape[-2])
        box_size_bias = torch.log(box_size + 1e-4) - torch.log1p(-box_size + 1e-4)

        # Compute box bias
        box_bias = torch.cat([box_coord_bias, box_size_bias], dim=-1)
        return box_bias

    def box_predictor(
        self,
        image_feats: torch.FloatTensor,
        feature_map: torch.FloatTensor,
    ) -> torch.FloatTensor:
        """
        Args:
            image_feats:
                Features extracted from the image, returned by the `image_text_embedder` method.
            feature_map:
                A spatial re-arrangement of image_features, also returned by the `image_text_embedder` method.
        Returns:
            pred_boxes:
                List of predicted boxes (cxcywh normalized to 0, 1) nested within a dictionary.
        """
        # Bounding box detection head [batch_size, num_boxes, 4].
        pred_boxes = self.box_head(image_feats)

        # Compute the location of each token on the grid and use it to compute a bias for the bbox prediction
        pred_boxes += self.compute_box_bias(feature_map)
        pred_boxes = self.sigmoid(pred_boxes)
        return pred_boxes

    def class_predictor(
        self,
        image_feats: torch.FloatTensor,
        query_embeds: Optional[torch.FloatTensor] = None,
        query_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.FloatTensor]:
        """
        Args:
            image_feats:
                Features extracted from the `image_text_embedder`.
            query_embeds:
                Text query embeddings.
            query_mask:
                Must be provided with query_embeddings. A mask indicating which query embeddings are valid.
        """
        (pred_logits, image_class_embeds) = self.class_head(image_feats, query_embeds, query_mask)

        return (pred_logits, image_class_embeds, query_mask)

    def forward(self, image_inputs, text_inputs_parts, text_embeds, targets):

        # Embed images and text queries
        input_ids = text_inputs_parts["input_ids"]
        text_embeds_parts, feature_map, outputs = self.image_text_embedder(input_ids=input_ids, attention_mask=text_inputs_parts["attention_mask"],
                                                                           pixel_values=image_inputs['pixel_values'])

        batch_size, num_patches, num_patches, hidden_dim = feature_map.shape
        image_feats = torch.reshape(feature_map, (batch_size, num_patches * num_patches, hidden_dim))

        # Reshape from [batch_size * max_text_queries, hidden_dim] -> [batch_size, max_text_queries, hidden_dim]
        max_text_queries = input_ids.shape[0] // batch_size
        text_embeds_parts = text_embeds_parts.reshape(batch_size, max_text_queries, text_embeds_parts.shape[-1])

        # If first token is 0, then this is a padded query [batch_size, num_queries].
        input_ids = input_ids.reshape(batch_size, max_text_queries, input_ids.shape[-1])
        query_mask = input_ids[..., 0] > 0

        # Store outputs for computing losses
        loss_dict = {}
        teacher_boxes_logits = torch.stack([target["logits"] for target in targets], dim=0)     # (bs, num_patches*num_patches, max_text_queries)

        if self.logits_from_teacher:
            pred_logits_parts = None
            topk_scores, topk_idxs = torch.topk(teacher_boxes_logits, k=1, dim=1)
        else:
            # Predict object classes [batch_size, num_patches, num_queries+1]
            pred_logits_parts, class_embeds, _ = self.class_predictor(image_feats, text_embeds_parts, query_mask)

            # Get the top-1 predictions
            scores = self.sigmoid(pred_logits_parts)
            topk_scores, topk_idxs = torch.topk(scores, k=1, dim=1)

        # Predict object boxes (cxcywh) --> Select 12 part boxes out of 576 predicted boxes
        pred_boxes = self.box_predictor(image_feats, feature_map)   # (bs, num_patches*num_patches, 4)
        pred_boxes_selected = torch.gather(pred_boxes, dim=1, index=topk_idxs.squeeze(1).unsqueeze(2).expand(*topk_idxs.squeeze(1).size(), pred_boxes.size(2)))

        # ----------------------------------------------------------------------------------------
        #   Computing box + class (unnecessary) + symmetric losses for box selection
        # ----------------------------------------------------------------------------------------
        outputs_loss = {}
        outputs_loss["logits"] = teacher_boxes_logits if self.logits_from_teacher else pred_logits_parts    # Not necessary as of August 26, 2023
        outputs_loss["pred_boxes"] = pred_boxes

        # Compute box + class losses
        if self.weight_dict["loss_ce"] > 0 or self.weight_dict["loss_bbox"] > 0 or self.weight_dict["loss_giou"] > 0:
            mapping_indices = [(selected_indices, torch.tensor(list(range(self.num_parts))).to(self.device)) for selected_indices in topk_idxs.squeeze(1)]
            loss_dict = self.criterion(outputs_loss, targets, mapping_indices)

        # For getting rid of the teacher model
        if self.weight_dict["loss_sym_box_label"] > 0 and not self.logits_from_teacher:
            # Compute symmetric loss to get rid of the teacher model
            logits_per_image = torch.softmax(pred_logits_parts, dim=1)
            logits_per_text = torch.softmax(pred_logits_parts, dim=-1)

            sym_loss_box_label = self.loss_symmetric(logits_per_image, logits_per_text, teacher_boxes_logits)
            loss_dict["loss_sym_box_label"] = sym_loss_box_label
        # ----------------------------------------------------------------------------------------

        # Predict image-level classes (batch_size, num_patches, num_queries)
        image_text_logits, pred_logits = self.cls_head(image_feats, text_embeds, topk_idxs)
        targets_cls = torch.tensor([target["targets_cls"] for target in targets]).unsqueeze(1).to(self.device)

        if self.weight_dict["loss_xclip"] > 0:
            if self.network_type == "classification":
                one_hot = torch.zeros_like(pred_logits).scatter(1, targets_cls, 1).to(self.device)

                # Compute CE loss
                if self.classification_loss == "ce_loss":
                    loss = self.ce_loss(pred_logits, one_hot)

                # OR Focal loss
                elif self.classification_loss == "focal_loss":
                    loss = self.focal_loss(pred_logits, targets_cls.squeeze(-1))

                    # Alternative implementation for the focal loss
                    # CE = F.cross_entropy(input=pred_logits, target=one_hot, reduction="none")
                    # prob = torch.gather(torch.softmax(input=pred_logits, dim=-1), 1, targets_cls.long()).squeeze(-1)
                    # loss = (self.alpha * ((1 - prob) ** self.gamma) * CE).mean()

                else:
                    raise f"Loss {self.classification_loss} is not supported"

                loss_dict["loss_xclip"] = loss
            else:
                # TODO: Need a linear classifier for this approach
                # Compute symmetric loss for part-descriptor contrastive learning
                logits_per_image = torch.softmax(image_text_logits, dim=0)
                logits_per_text = torch.softmax(image_text_logits, dim=-1)
                sym_loss = self.loss_symmetric(logits_per_image, logits_per_text, targets_cls)
                loss_dict["loss_xclip"] = sym_loss

        return pred_logits, center_to_corners_format(pred_boxes_selected), loss_dict

    def loss_symmetric(self, text_logits: torch.Tensor, image_logits: torch.Tensor, targets: torch.Tensor, box_labels: torch.Tensor = None) -> torch.Tensor:
        # text/image logits (batch_size*num_boxes, num_classes*num_descs): The logits that softmax over text descriptors or boxes
        # targets (batch_size, 1): The ground truth label of box-text pair for classification OR
        # targets (batch_size, all_boxes, num_parts): The ground truth label of box-text pair for box selection
        # box_labels (batch_size, num_boxes), 0 for no box, 1 for box

        assert text_logits.shape == image_logits.shape

        # For image classification
        if image_logits.shape != targets.shape:
            batch_size = targets.shape[0]

            # get the matching labels (bs * 12, num_classes * num_parts)
            default_box_labels = torch.kron(torch.ones(batch_size, self.num_classes), torch.eye(self.num_parts)).to(self.device)
            if box_labels is None:
                box_labels = default_box_labels.clone()
            else:
                # (batch_size, num_boxes) -> (bs * num_boxes, num_classes * num_parts)
                box_labels = box_labels.view(-1, 1) * default_box_labels

            # Create one-hot encoding of targets; matching_labels shape: (bs * 12, num_classes * num_parts)
            target_one_hot = torch.zeros(batch_size, self.num_classes).to(self.device).scatter(1, targets.view(-1, 1), 1)
            target_one_hot = torch.kron(target_one_hot, torch.ones(self.num_parts, self.num_parts).to(self.device))

            matching_labels = target_one_hot * box_labels
        else:
            # For box selection: matching_labels shape: (bs, 576, num_parts)
            values, indices = torch.max(targets, dim=1)
            matching_labels = torch.zeros_like(targets).scatter(1, indices.unsqueeze(1), 1)

        loss_i = F.cross_entropy(image_logits, matching_labels, reduction='mean')
        loss_t = F.cross_entropy(text_logits, matching_labels, reduction='mean')
        sym_loss = (loss_i + loss_t).mean()

        return sym_loss


class FocalLoss(nn.Module):
    def __init__(self, alpha=1.0, gamma=2.0, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, input, target):
        # Compute the Softmax and Cross-Entropy Loss
        ce_loss = F.cross_entropy(input, target, reduction='none')

        # Convert the target labels to one-hot format
        num_classes = input.size(1)  # Number of classes
        target_onehot = torch.zeros_like(input)
        target_onehot.scatter_(1, target.unsqueeze(1), 1)

        # Compute the probabilities of the true labels
        prob = F.softmax(input, dim=1)
        true_prob = torch.sum(prob * target_onehot, dim=1)

        # Compute the Focal Loss
        focal_weight = (self.alpha * ((1 - true_prob) ** self.gamma))
        focal_loss = focal_weight * ce_loss

        # Apply the reduction method (mean, sum, or none)
        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss


class DetrHungarianMatcher(nn.Module):
    """
    This class computes an assignment between the targets and the predictions of the network.

    For efficiency reasons, the targets don't include the no_object. Because of this, in general, there are more
    predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, while the others are
    un-matched (and thus treated as non-objects).

    Args:
        class_cost:
            The relative weight of the classification error in the matching cost.
        bbox_cost:
            The relative weight of the L1 error of the bounding box coordinates in the matching cost.
        giou_cost:
            The relative weight of the giou loss of the bounding box in the matching cost.
    """

    def __init__(self, class_cost: float = 1, bbox_cost: float = 1, giou_cost: float = 1):
        super().__init__()
        requires_backends(self, ["scipy"])

        self.class_cost = class_cost
        self.bbox_cost = bbox_cost
        self.giou_cost = giou_cost
        if class_cost == 0 and bbox_cost == 0 and giou_cost == 0:
            raise ValueError("All costs of the Matcher can't be 0")

    @torch.no_grad()
    def forward(self, outputs, targets):
        """
        Args:
            outputs (`dict`):
                A dictionary that contains at least these entries:
                * "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
                * "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates.
            targets (`List[dict]`):
                A list of targets (len(targets) = batch_size), where each target is a dict containing:
                * "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of
                  ground-truth
                 objects in the target) containing the class labels
                * "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates.

        Returns:
            `List[Tuple]`: A list of size `batch_size`, containing tuples of (index_i, index_j) where:
            - index_i is the indices of the selected predictions (in order)
            - index_j is the indices of the corresponding selected targets (in order)
            For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
        """
        batch_size, num_queries = outputs["logits"].shape[:2]

        # We flatten to compute the cost matrices in a batch
        out_prob = outputs["logits"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]
        out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]

        # Also concat the target labels and boxes
        target_ids = torch.cat([v["class_labels"] for v in targets])
        target_bbox = torch.cat([v["boxes"] for v in targets])

        # Compute the classification cost. Contrary to the loss, we don't use the NLL,
        # but approximate it in 1 - proba[target class].
        # The 1 is a constant that doesn't change the matching, it can be ommitted.
        class_cost = -out_prob[:, target_ids]

        # Compute the L1 cost between boxes
        bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)

        # Compute the giou cost between boxes
        giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))

        # Final cost matrix
        cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
        cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()

        sizes = [len(v["boxes"]) for v in targets]
        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]
        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]


class DetrLoss(nn.Module):
    """
    This class computes the losses for DetrForObjectDetection/DetrForSegmentation. The process happens in two steps: 1)
    we compute hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair
    of matched ground-truth / prediction (supervise class and box).

    A note on the `num_classes` argument (copied from original repo in detr.py): "the naming of the `num_classes`
    parameter of the criterion is somewhat misleading. It indeed corresponds to `max_obj_id` + 1, where `max_obj_id` is
    the maximum id for a class in your dataset. For example, COCO has a `max_obj_id` of 90, so we pass `num_classes` to
    be 91. As another example, for a dataset that has a single class with `id` 1, you should pass `num_classes` to be 2
    (`max_obj_id` + 1). For more details on this, check the following discussion
    https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223"


    Args:
        matcher (`DetrHungarianMatcher`):
            Module able to compute a matching between targets and proposals.
        num_parts (`int`):
            Number of object categories, omitting the special no-object category.
        eos_coef (`float`):
            Relative classification weight applied to the no-object category.
        losses (`List[str]`):
            List of all the losses to be applied. See `get_loss` for a list of all available losses.
    """

    def __init__(self, matcher, num_parts, eos_coef, losses):
        super().__init__()
        self.matcher = matcher
        self.num_parts = num_parts
        self.eos_coef = eos_coef
        self.losses = losses

        # empty_weight = torch.ones(self.num_parts + 1)
        empty_weight = torch.ones(self.num_parts)
        empty_weight[-1] = self.eos_coef
        self.register_buffer("empty_weight", empty_weight)

    # NOTE: WE DO NOT USE BOX LABEL LOSS IN OUR METHOD
    # Removed logging parameter, which was part of the original implementation
    def loss_labels(self, outputs, targets, indices, num_boxes):
        """
        Classification loss (NLL) targets dicts must contain the key "class_labels" containing a tensor of dim
        [nb_target_boxes]
        """
        if "logits" not in outputs:
            raise KeyError("No logits were found in the outputs")
        source_logits = outputs["logits"]

        idx = self._get_source_permutation_idx(indices)
        # target_classes_o = torch.cat([t["class_labels"][J] for t, (_, J) in zip(targets, indices)])
        # target_classes = torch.full(source_logits.shape[:2], self.num_parts, dtype=torch.int64, device=source_logits.device)
        # target_classes[idx] = target_classes_o

        source_logits = source_logits[idx].view(len(indices), -1, self.num_parts)
        target_classes = torch.stack([t["class_labels"][J] for t, (_, J) in zip(targets, indices)], dim=0)

        loss_ce = F.cross_entropy(source_logits.transpose(1, 2), target_classes, self.empty_weight)
        losses = {"loss_ce": loss_ce}

        return losses

    @torch.no_grad()
    def loss_cardinality(self, outputs, targets, indices, num_boxes):
        """
        Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes.

        This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients.
        """
        logits = outputs["logits"]
        device = logits.device
        target_lengths = torch.as_tensor([len(v["class_labels"]) for v in targets], device=device)
        # Count the number of predictions that are NOT "no-object" (which is the last class)
        card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1)
        card_err = nn.functional.l1_loss(card_pred.float(), target_lengths.float())
        losses = {"cardinality_error": card_err}
        return losses

    def loss_boxes(self, outputs, targets, indices, num_boxes):
        """
        Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss.

        Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes
        are expected in format (center_x, center_y, w, h), normalized by the image size.
        """
        if "pred_boxes" not in outputs:
            raise KeyError("No predicted boxes found in outputs")

        device = outputs["pred_boxes"].device

        idx = self._get_source_permutation_idx(indices)
        # source_boxes = center_to_corners_format(outputs["pred_boxes"][idx])                         # corner format
        # target_boxes = torch.cat([torch.tensor(t["boxes"]).to(device) for t in targets], dim=0)     # corner format

        source_boxes = outputs["pred_boxes"][idx]                                                     # center format
        target_boxes = torch.cat([torch.tensor(t["boxes"]).to(device) for t in targets], dim=0)       # corner format
        target_boxes = corners_to_center_format(target_boxes)                                         # center format

        losses = {}

        loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="mean")
        losses["loss_bbox"] = loss_bbox

        # Since we have part-part mappings already, the giou loss for each part is from the diagonal of the gIoU matrix
        loss_giou = 1 - torch.diag(generalized_box_iou(center_to_corners_format(source_boxes), center_to_corners_format(target_boxes)))
        losses["loss_giou"] = loss_giou.mean()

        return losses

    def loss_masks(self, outputs, targets, indices, num_boxes):
        """
        Compute the losses related to the masks: the focal loss and the dice loss.

        Targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w].
        """
        if "pred_masks" not in outputs:
            raise KeyError("No predicted masks found in outputs")

        source_idx = self._get_source_permutation_idx(indices)
        target_idx = self._get_target_permutation_idx(indices)
        source_masks = outputs["pred_masks"]
        source_masks = source_masks[source_idx]
        masks = [t["masks"] for t in targets]

        # TODO use valid to mask invalid areas due to padding in loss
        target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
        target_masks = target_masks.to(source_masks)
        target_masks = target_masks[target_idx]

        # upsample predictions to the target size
        source_masks = nn.functional.interpolate(
            source_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False
        )
        source_masks = source_masks[:, 0].flatten(1)

        target_masks = target_masks.flatten(1)
        target_masks = target_masks.view(source_masks.shape)
        losses = {
            "loss_mask": sigmoid_focal_loss(source_masks, target_masks, num_boxes),
            "loss_dice": dice_loss(source_masks, target_masks, num_boxes),
        }
        return losses

    def _get_source_permutation_idx(self, indices):
        # permute predictions following indices
        batch_idx = torch.cat([torch.full_like(source, i) for i, (source, _) in enumerate(indices)])
        source_idx = torch.cat([source for (source, _) in indices])
        return batch_idx, source_idx

    def _get_target_permutation_idx(self, indices):
        # permute targets following indices
        batch_idx = torch.cat([torch.full_like(target, i) for i, (_, target) in enumerate(indices)])
        target_idx = torch.cat([target for (_, target) in indices])
        return batch_idx, target_idx

    def get_loss(self, loss, outputs, targets, indices, num_boxes):
        loss_map = {
            "labels": self.loss_labels,
            "cardinality": self.loss_cardinality,
            "boxes": self.loss_boxes,
            "masks": self.loss_masks,
        }
        if loss not in loss_map:
            raise ValueError(f"Loss {loss} not supported")
        return loss_map[loss](outputs, targets, indices, num_boxes)

    def forward(self, outputs, targets, indices):
        """
        This performs the loss computation.

        Args:
             outputs (`dict`, *optional*):
                Dictionary of tensors, see the output specification of the model for the format.
             targets (`List[dict]`, *optional*):
                List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the
                losses applied, see each loss' doc.
        """
        outputs_without_aux = {k: v for k, v in outputs.items() if k != "auxiliary_outputs"}

        # ThangPM: Do NOT use bipartite matching --> Use the boxes selected by argmax for computing symmetric loss
        # Retrieve the matching between the outputs of the last layer and the targets
        # indices = self.matcher(outputs_without_aux, targets)

        # Compute the average number of target boxes across all nodes, for normalization purposes
        num_boxes = sum(len(t["class_labels"]) for t in targets)
        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=outputs["pred_boxes"].device)
        # (Niels): comment out function below, distributed training to be added
        # if is_dist_avail_and_initialized():
        #     torch.distributed.all_reduce(num_boxes)
        # (Niels) in original implementation, num_boxes is divided by get_world_size()
        num_boxes = torch.clamp(num_boxes, min=1).item()

        # Compute all the requested losses
        losses = {}
        for loss in self.losses:
            losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))

        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if "auxiliary_outputs" in outputs:
            for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]):
                # indices = self.matcher(auxiliary_outputs, targets)
                for loss in self.losses:
                    if loss == "masks":
                        # Intermediate masks losses are too costly to compute, we ignore them.
                        continue
                    l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes)
                    l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
                    losses.update(l_dict)

        return losses


class NestedTensor(object):
    def __init__(self, tensors, mask: Optional[Tensor]):
        self.tensors = tensors
        self.mask = mask

    def to(self, device):
        cast_tensor = self.tensors.to(device)
        mask = self.mask
        if mask is not None:
            cast_mask = mask.to(device)
        else:
            cast_mask = None
        return NestedTensor(cast_tensor, cast_mask)

    def decompose(self):
        return self.tensors, self.mask

    def __repr__(self):
        return str(self.tensors)


def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
    def _max_by_axis(the_list):
        # type: (List[List[int]]) -> List[int]
        maxes = the_list[0]
        for sublist in the_list[1:]:
            for index, item in enumerate(sublist):
                maxes[index] = max(maxes[index], item)
        return maxes

    if tensor_list[0].ndim == 3:
        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
        batch_shape = [len(tensor_list)] + max_size
        batch_size, num_channels, height, width = batch_shape
        dtype = tensor_list[0].dtype
        device = tensor_list[0].device
        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
        mask = torch.ones((batch_size, height, width), dtype=torch.bool, device=device)
        for img, pad_img, m in zip(tensor_list, tensor, mask):
            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
            m[: img.shape[1], : img.shape[2]] = False
    else:
        raise ValueError("Only 3-dimensional tensors are supported")
    return NestedTensor(tensor, mask)


def dice_loss(inputs, targets, num_boxes):
    """
    Compute the DICE loss, similar to generalized IOU for masks

    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs (0 for the negative class and 1 for the positive
                 class).
    """
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1)
    numerator = 2 * (inputs * targets).sum(1)
    denominator = inputs.sum(-1) + targets.sum(-1)
    loss = 1 - (numerator + 1) / (denominator + 1)
    return loss.sum() / num_boxes


def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2):
    """
    Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.

    Args:
        inputs (`torch.FloatTensor` of arbitrary shape):
            The predictions for each example.
        targets (`torch.FloatTensor` with the same shape as `inputs`)
            A tensor storing the binary classification label for each element in the `inputs` (0 for the negative class
            and 1 for the positive class).
        alpha (`float`, *optional*, defaults to `0.25`):
            Optional weighting factor in the range (0,1) to balance positive vs. negative examples.
        gamma (`int`, *optional*, defaults to `2`):
            Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples.

    Returns:
        Loss tensor
    """
    prob = inputs.sigmoid()
    ce_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
    # add modulating factor
    p_t = prob * targets + (1 - prob) * (1 - targets)
    loss = ce_loss * ((1 - p_t) ** gamma)

    if alpha >= 0:
        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
        loss = alpha_t * loss

    return loss.mean(1).sum() / num_boxes


src/owlvit_cls.py
anguyen8/peeb
import os
import copy
import shutil
import argparse
from typing import Dict, List, Optional, Tuple, Union

import torch
import numpy as np
from PIL import Image
from tqdm import tqdm
from torch.utils.data import DataLoader
from transformers import OwlViTProcessor, OwlViTForObjectDetection
from torchvision.transforms import Resize, InterpolationMode, ToPILImage

from data_loader import BirdSoup


def xy_in_box(box: list[int], x: int, y: int):
    x1, y1, x2, y2 = box
    return x1 <= x <= x2 and y1 <= y <= y2


def check_exist(image_ids: list[str], box_dir: str, ext: str = ".pth") -> tuple[bool, list[str]]:
    missing_ids = [
        image_id
        for image_id in image_ids
        if not os.path.exists(os.path.join(box_dir, image_id + ext))
    ]
    all_exist = not missing_ids
    return all_exist, missing_ids


def bound_boxes(boxes: torch.Tensor, img_sizes: torch.Tensor):
    # note: img_sizes are in (h, w) format
    
    # Create a tensor of boundaries based on image sizes
    boundaries = torch.cat((torch.zeros_like(img_sizes), img_sizes.flip(1)), dim=1)
    # Bounding the boxes
    bounded_boxes = torch.zeros_like(boxes)
    bounded_boxes[:, 0] = torch.clamp(boxes[:, 0], boundaries[:, 0], boundaries[:, 2])
    bounded_boxes[:, 1] = torch.clamp(boxes[:, 1], boundaries[:, 1], boundaries[:, 3])
    bounded_boxes[:, 2] = torch.clamp(boxes[:, 2], boundaries[:, 0], boundaries[:, 2])
    bounded_boxes[:, 3] = torch.clamp(boxes[:, 3], boundaries[:, 1], boundaries[:, 3])

    return bounded_boxes.long()


# to return PIL images
def custom_collate_fn(batch):
    images, targets, paths, image_sizes = zip(*batch)
    targets = torch.tensor(targets, dtype=torch.long)
    images_sizes = torch.stack(image_sizes)
    return images, targets, paths, images_sizes


# sourcery skip: assign-if-exp, swap-if-expression
if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    # OwlViT settings
    parser.add_argument('--owl_model', help='select owl-vit version (huggingface)', default="owlvit-large-patch14", choices=["owlvit-base-patch32", "owlvit-base-patch16", "owlvit-large-patch14"])
    parser.add_argument('--owl_prompt_type', type=str, help='Prompt type for OwlViT', default="cub-12-parts")

    # Misc
    parser.add_argument('--meta_path', type=str, help='Path to meta file', default="../data/bird_11K/metadata/bird_soup_uncased_v2.h5")
    parser.add_argument('--batch_size', type=int, help='num batch size', default=12)
    parser.add_argument('--num_workers', type=int, help='num workers for batch processing', default=8)
    parser.add_argument('--device', help='select device', default="cuda:0", type=str)
    parser.add_argument('--overwrite', help='overwrite existing boxes results', action='store_true')
    parser.add_argument('--output_subfix', type=str, help='output subfix', default=None)
    parser.add_argument('--filter_size', type=int, default=None, help='Filter size for bird, default is None (no filter)')
    parser.add_argument('--save_img', action='store_true', help='Save the image to a new folder')
    parser.add_argument('--image_output_dir', type=str, help='Output directory', default="../data/bird_11K/images")
    parser.add_argument('--owl_threshold', type=float, help='OwlViT threshold, recommand set to 0, so we can filter later.', default=0.0)
    parser.add_argument('--shuffle', action='store_true', help='Shuffle the dataset')
    parser.add_argument('--save_logits', action='store_true', help='Save logits')

    args = parser.parse_args()

    if args.owl_prompt_type == "cub-12-parts":
        prompts = ["bird", "back", "beak", "belly", "breast", "crown", "forehead", "eyes", "legs", "wings", "nape", "tail", "throat"]
    else:
        raise NotImplementedError(f"Prompt type {args.owl_prompt_type} is not implemented")

    boxes_dir = f"../data/bird_11K/part_boxes/{args.owl_model}_{args.owl_prompt_type}"
    if args.output_subfix is not None:
        boxes_dir = f"{boxes_dir}_{args.output_subfix}"

    os.makedirs(boxes_dir, exist_ok=True)
    os.makedirs(args.image_output_dir, exist_ok=True)
    device = args.device

    # Load model
    processor = OwlViTProcessor.from_pretrained(f"google/{args.owl_model}")
    model = OwlViTForObjectDetection.from_pretrained(f"google/{args.owl_model}")
    model.to(device)
    model.eval()

    # Load dataset and prepare data loader
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    dataset = BirdSoup(root=None, meta_path=args.meta_path, use_meta_dir=True, transform=processor)
    dataloader = DataLoader(dataset, args.batch_size, shuffle=args.shuffle, num_workers=args.num_workers, pin_memory=True)

    # Prepare text embeddings
    idx2name = dataset.idx2class

    with torch.no_grad():
        # tokenize descs
        text_input_tokens = processor(prompts, padding="max_length", truncation=True, return_tensors="pt")
        text_input_tokens['input_ids'] = text_input_tokens['input_ids'].repeat(args.batch_size, 1)
        text_input_tokens['attention_mask'] = text_input_tokens['attention_mask'].repeat(args.batch_size, 1)
        text_input_tokens = text_input_tokens.to(device)

        for batch_idx, batch in tqdm(enumerate(dataloader), desc='Computing boxes', total=len(dataloader)):
            images, class_labels, img_sizes, paths, soup_names = batch

            image_ids = []
            output_image_paths = []
            for image_name in soup_names:
                image_ids.append(os.path.splitext(image_name)[0])
                output_image_paths.append(os.path.join(args.image_output_dir, image_name))

            # check if boxes already exist
            if not args.overwrite:
                all_exist, missing_ids = check_exist(image_ids, boxes_dir)
                if all_exist:
                    continue  
                else:
                    keep_missing = [id in missing_ids for id in image_ids]
            else:
                keep_missing = [True] * len(image_ids)

            # forward to get owl-vit outputs
            images['pixel_values'] = images['pixel_values'].squeeze(1).to(device)
            owl_inputs = images | text_input_tokens
            owl_outputs = model(**owl_inputs)
            preds = processor.post_process_object_detection(outputs=owl_outputs, threshold=0, target_sizes=img_sizes.to(device))
            all_boxes = torch.stack([pred['boxes'] for pred in preds]).cpu()
            logits = owl_outputs.logits
            scores_ = torch.sigmoid(logits).cpu()

            top1_scores, top1_idxs = torch.topk(scores_, k=1, dim=1)

            # derive the boxes of the bird
            bird_box_lower_bound = args.filter_size ** 2 if args.filter_size is not None else 0
            bird_scores = top1_scores[:, 0, 0]
            bird_idxs = top1_idxs[:, 0, 0]

            # filter out bird boxes that are too small
            bird_boxes = all_boxes[torch.arange(len(all_boxes)), bird_idxs]

            # calculate the area of the bird boxes, and remove birds that are too small
            areas = (bird_boxes[:, 2] - bird_boxes[:, 0]) * (bird_boxes[:, 3] - bird_boxes[:, 1])

            if args.filter_size is not None:
                score_keep = bird_scores > args.owl_threshold
                box_keep = areas > bird_box_lower_bound
                keep = score_keep & box_keep & torch.tensor(keep_missing)
            else:
                keep = [True] * len(all_boxes)

            part_scores = top1_scores[:, 0, 1:]
            part_idxs = top1_idxs[:, 0, 1:]

            # filter by keep
            part_scores = part_scores[keep]
            part_idxs = part_idxs[keep]
            bird_scores = bird_scores[keep]
            bird_boxes = bird_boxes[keep]
            part_logits = scores_[keep][:, :, 1:].cpu()

            # choose boxes per query and store class-wise per image
            for i, (boxes, part_idx, part_score, part_logit, bird_box, bird_score) in enumerate(zip(all_boxes, part_idxs, part_scores, part_logits, bird_boxes, bird_scores)):
                
                part_boxes = boxes[part_idx]
                part_boxes = bound_boxes(boxes=part_boxes, img_sizes=img_sizes[i].repeat(len(part_boxes), 1))

                img_pred_dict = {
                    part_name: box.cpu().tolist()
                    for part_name, box in zip(prompts[1:], part_boxes)
                }

                torch.save({"image_id": image_ids[i],
                            "image_path": output_image_paths[i],
                            "org_image_path": paths[i],
                            "image_size": img_sizes[i].tolist()[::-1], # in (w, h)
                            "class_name": idx2name[class_labels[i].item()],
                            "boxes_info": img_pred_dict, # previous key should be 'part_boxes', keep it for consistency for Thang's code
                            "part_scores": part_score.tolist(),
                            "part_logits": part_logit if args.save_logits else None,
                            
                            "bird_boxes": bird_box.tolist(),
                            "bird_score": bird_score.item(),
                            "bird_areas": areas[i].item(),
                            "bird_ratio": (areas[i].item() / (img_sizes[i][0] * img_sizes[i][1])).item(),
                            }, f"{boxes_dir}/{image_ids[i]}.pth")
                
                if args.save_img:
                    shutil.copy(paths[i], output_image_paths[i])




src/precompute_boxes.py
anguyen8/peeb
import os
import shutil
import argparse

import pandas as pd
from tqdm import tqdm


def get_arg():
    argparser = argparse.ArgumentParser()
    argparser.add_argument("--meta_path", type=str)
    argparser.add_argument("--img_path", type=str)
    
    return argparser.parse_args()

def rename_images(meta_path: str, img_path: str, inplace: bool = True):
    if not inplace:
        # add suffix to image folder
        out_image_path = f"{img_path}_birdsoup"
    else:
        out_image_path = img_path
    os.makedirs(out_image_path, exist_ok=True)    
        
    meta = pd.read_hdf(meta_path)

    missing_images = []
    for idx, row in tqdm(meta.iterrows(), total=len(meta), desc="Rename images"):
        org_name = row["org_image_name"]
        new_image_name = row["image_name"]

        org_image_path = os.path.join(img_path, org_name)
        new_image_path = os.path.join(out_image_path, new_image_name)
        # skip and record missing images
        if not os.path.exists(org_image_path):
            data_source = row["data_source"]
            missing_images.append((org_image_path, data_source))
            continue
        # skip if image already exists (in case we want to run this script multiple times)
        if os.path.exists(new_image_path):
            continue
        
        shutil.move(org_image_path, new_image_path)
        
if __name__ == "__main__":
    args = get_arg()
    rename_images(args.meta_path, args.img_path, inplace=True)
src/rename_birdsoup_images.py
anguyen8/peeb
import argparse
import gc
import json
import math
import os.path
import statistics
from datetime import datetime

import pytz
import wandb

import numpy as np
import spacy
import torchvision
import torchmetrics
from sklearn.model_selection import train_test_split
import torch
from torch.utils.data import DataLoader, Dataset, Subset
from torch import nn
from tqdm import tqdm
from transformers import OwlViTProcessor, OwlViTForObjectDetection
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

from data_loader import CUB, NABirdsDataset, BirdSoup
import torch.multiprocessing as mp

from functools import reduce
import operator
import pandas as pd

from PIL import Image, ImageDraw, ImageFont
from transformers.image_transforms import corners_to_center_format
from transformers.models.detr.modeling_detr import generalized_box_iou, center_to_corners_format, box_iou

from src.owlvit_cls_plus import OwlViTForClassification

from utils import *

nlp_spacy = spacy.load("en_core_web_sm")


def get_timestamp():
    local_tz = pytz.timezone("America/Chicago")
    return datetime.now(tz=local_tz).strftime("%Y%m%d_%H%M%S")


def seed_everything(random_state: int):
    random.seed(random_state)
    np.random.seed(random_state)
    torch.manual_seed(random_state)
    torch.cuda.manual_seed(random_state)
    if random_state == 0:  # slower, more reproducible
        torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic = False, True
    else:  # faster, less reproducible
        torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic = True, False


def check_device_availability(devices: list[int]):
    # get all available devices
    available_devices = torch.cuda.device_count()
    available_list = list(range(available_devices))
    # check if the device is available
    for device in devices:
        if int(device) not in available_list:
            raise ValueError(f"Device {device} is not available. Available devices are {available_list}")


# get a train, val, and test split
def get_data_split(dataset: Dataset, labels: list = None, eval_size: float = 0.2, random_state: int = 42):
    if labels is None:
        try:
            labels = dataset.labels
        except:
            raise ValueError("Labels are not provided and the dataset does not have a labels attribute")
    if eval_size == 0:
        train_split = dataset
        val_split = None
    else:
        train_indices, evaluate_indices, _, _ = train_test_split(range(len(dataset)), labels, stratify=labels, test_size=eval_size, random_state=random_state)
        train_split = Subset(dataset, train_indices)
        val_split = Subset(dataset, evaluate_indices)

    return train_split, val_split


def load_training_dataset(dataset_name: str, sub_dataset_names: str, eval_size: float, transform: dict[str: torchvision.transforms.Compose] or torchvision.transforms.Compose, random_state: int = 42, split: str = 'train', zeroshot_split: bool = False):
    """_summary_

    Args:
        dataset_name (str): name of dataset for training
        eval_size (float): proportion of validation set extracted from the train set. Default to 0.2
        transform (dict): dict[str: torchvision.transforms.Compose], the key is the name of the transform, the value is the transform module
        random_state (int, optional): random seed for spliting train and val when eval_size is not 0. Defaults to 42.
    """
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    if dataset_name == 'cub':
        if split != 'test':
            dataset_org = CUB(CUB_DIR, transform=transform, train=True, return_path=True, zeroshot_split=zeroshot_split)
            dataset, val_dataset = get_data_split(dataset_org, labels=dataset_org.targets, eval_size=eval_size, random_state=random_state)
            val_dataset.transform = transform
        else:
            dataset = CUB(CUB_DIR, transform=transform, train=False, return_path=True, zeroshot_split=zeroshot_split)
            val_dataset = None

    elif dataset_name == "nabirds":
        if split != 'test':
            dataset_org = NABirdsDataset(NABIRDS_DIR, transform=transform, train=True, return_path=True, zeroshot_split=zeroshot_split)
            dataset, val_dataset = get_data_split(dataset_org, labels=dataset_org.targets, eval_size=eval_size, random_state=random_state)
            val_dataset.transform = transform
        else:
            dataset = NABirdsDataset(NABIRDS_DIR, transform=transform, train=False, return_path=True, zeroshot_split=zeroshot_split)
            val_dataset = None

    elif dataset_name == "bird_soup":
        sub_datasets = [sub_dataset.strip() for sub_dataset in sub_dataset_names.split(',')] if sub_dataset_names is not None else None

        if split != 'test':
            dataset = BirdSoup(BIRD_SOUP_DIR, transform=transform, train=True, return_path=True, meta_path=args.train_file, subset=sub_datasets)
            val_dataset = BirdSoup(BIRD_SOUP_DIR, transform=transform, train=False, return_path=True, meta_path=args.val_file, subset=sub_datasets)
        else:
            dataset = BirdSoup(BIRD_SOUP_DIR, transform=transform, train=False, return_path=True, meta_path=args.test_file, subset=sub_datasets)
            val_dataset = None

    return dataset, val_dataset


def visualize_bbox(image_path, gt_bboxes, 
                   base_bboxes, base_bbox_losses, base_giou_losses, base_iou_scores, 
                   pred_bboxes, bbox_losses, giou_losses, iou_scores, store_path=None):
    """
    Draw bounding boxes on the image
    :param image_path: path to image
    :param gt_bboxes: tensor of ground truth bounding boxes. shape (12, 4)
    :param pred_bboxes: tensor of predicted bounding boxes. shape (12, 4)
    :returns: Image object with bounding boxes drawn on it
    """

    def completely_overlapped(box1, box2):
        """Check if box1 is completely overlapped by box2."""
        return box2[0] <= box1[0] and box2[1] <= box1[1] and box2[2] >= box1[2] and box2[3] >= box1[3]

    # Open images
    image_gt = Image.open(image_path)
    draw_gt = ImageDraw.Draw(image_gt)

    image_base = Image.open(image_path)
    draw_base = ImageDraw.Draw(image_base)

    image_pred = Image.open(image_path)
    draw_pred = ImageDraw.Draw(image_pred)

    # Define a list of 12 colors
    colors = [
        (255, 0, 0), (0, 255, 0), (0, 0, 255),          # RGB colors: red, green, blue
        (255, 255, 0), (0, 255, 255), (255, 0, 255),    # RGB colors: yellow, cyan, magenta
        (128, 0, 0), (0, 128, 0), (0, 0, 128),          # Darker versions of above colors
        (128, 128, 0), (0, 128, 128), (128, 0, 128)     # Darker versions of above colors
    ]

    part_names = ['back', 'beak', 'belly', 'breast', 'crown', 'forehead', 'eyes', 'legs', 'wings', 'nape', 'tail', 'throat']

    font = ImageFont.load_default()
    base_legend_text = []
    legend_text = []

    # Detect overlap between boxes and, if one is completely overlapped, add margins
    for i in range(len(pred_bboxes)):
        for j in range(i + 1, len(pred_bboxes)):
            random_margin = np.random.uniform(0, 5, size=4)

            if completely_overlapped(pred_bboxes[i], pred_bboxes[j]):
                pred_bboxes[i] += random_margin

            if completely_overlapped(base_bboxes[i], base_bboxes[j]):
                base_bboxes[i] += random_margin

            if completely_overlapped(gt_bboxes[i], gt_bboxes[j]):
                gt_bboxes[i] += random_margin

    # Draw bounding boxes. Each box will be drawn with a different color.
    for box, color, part_name in zip(gt_bboxes, colors, part_names):
        draw_gt.rectangle(box.tolist(), outline=color, width=2)
        # draw_gt.text((box[0], box[1]), part_name, fill=color)

    for box, color, part_name, loss_bbox, loss_giou, iou_score in zip(base_bboxes, colors, part_names, base_bbox_losses, base_giou_losses, base_iou_scores):
        draw_base.rectangle(box.tolist(), outline=color, width=2)
        base_legend_text.append((part_name + "|" + str(round(loss_bbox.item(), 2)) + "|" + str(round(loss_giou.item(), 2)) + "|" + str(round(iou_score.item(), 2)), color))

    for box, color, part_name, loss_bbox, loss_giou, iou_score in zip(pred_bboxes, colors, part_names, bbox_losses, giou_losses, iou_scores):
        draw_pred.rectangle(box.tolist(), outline=color, width=2)
        legend_text.append((part_name + "|" + str(round(loss_bbox.item(), 2)) + "|" + str(round(loss_giou.item(), 2)) + "|" + str(round(iou_score.item(), 2)), color))

    base_legend_text.append((f"Average|{round(base_bbox_losses.mean().item(), 2)}|{round(base_giou_losses.mean().item() ,2)}|{round(base_iou_scores.mean().item() ,2)}", (255, 128, 0)))
    legend_text.append((f"Average|{round(bbox_losses.mean().item(), 2)}|{round(giou_losses.mean().item() ,2)}|{round(iou_scores.mean().item() ,2)}", (255, 128, 0)))

    # Draw legend texts
    for i, (text, color) in enumerate(base_legend_text):
        draw_base.text((0, i*10), text, fill=color, font=font)

    for i, (text, color) in enumerate(legend_text):
        draw_pred.text((0, i*10), text, fill=color, font=font)

    if store_path is not None:
        # Concatenate horizontally
        image = Image.new('RGB', (image_gt.width + image_base.width + image_pred.width, image_gt.height))
        image.paste(image_gt, (0, 0))
        image.paste(image_base, (image_gt.width, 0))
        image.paste(image_pred, (image_gt.width + image_base.width, 0))

        # Save the output image
        image.save(store_path + "/" + image_path.split("/")[-1].replace(".jpg", ".png"))


def forward_inputs(model, images, text_inputs_parts, text_embeds, targets, weight_dict, visualize_batch_count=0, store_path=None):
    '''
        image_embeds.shape = {Size} torch.Size([10, 60, 60, 1024])
        text_embeds.shape = {Size} torch.Size([10, 2400, 768])
        owlvit_logits.shape = {Size} torch.Size([10, 3600, 12])
    '''

    pred_logits, pred_boxes, loss_dict = model(images, text_inputs_parts, text_embeds, targets)

    # Compute total loss, as a weighted sum of the various losses
    loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)

    target_boxes = torch.stack([torch.tensor(t["boxes"]) for t in targets], dim=0)
    base_boxes = torch.stack([torch.tensor(t["boxes_base"]) for t in targets], dim=0)
    pred_boxes = pred_boxes.detach().cpu()
    
    # giou_scores = torch.diag(generalized_box_iou(pred_boxes.view(-1, 4), target_boxes.view(-1, 4))).view(-1, 12)
    iou_scores = torch.diag(box_iou(pred_boxes.view(-1, 4), target_boxes.view(-1, 4))[0]).view(-1, 12)
    base_iou_scores = torch.diag(box_iou(base_boxes.view(-1, 4), target_boxes.view(-1, 4))[0]).view(-1, 12)

    # loss_dict["giou_score"] = giou_scores.mean()
    loss_dict["iou_score"] = iou_scores.mean()
    loss_dict["base_iou_scores"] = base_iou_scores.mean()

    if visualize_batch_count > 0:
        image_paths = [t["image_path"] for t in targets]
        image_sizes = [t["image_size"] for t in targets]

        base_bbox_losses = torch.mean(nn.functional.l1_loss(corners_to_center_format(base_boxes), corners_to_center_format(target_boxes), reduction="none"), dim=-1)
        base_giou_losses = (1 - torch.diag(generalized_box_iou(base_boxes.view(-1, 4), target_boxes.view(-1, 4)))).view(-1, 12)

        bbox_losses = torch.mean(nn.functional.l1_loss(corners_to_center_format(pred_boxes), corners_to_center_format(target_boxes), reduction="none"), dim=-1)
        giou_losses = (1 - torch.diag(generalized_box_iou(pred_boxes.view(-1, 4), target_boxes.view(-1, 4)))).view(-1, 12)

        for image_path, image_size, gt_bbox, base_bbox, base_bbox_loss, base_giou_loss, base_iou_score, pred_bbox, bbox_loss, giou_loss, iou_score in zip(image_paths, image_sizes, target_boxes, 
                                                                                                                                                          base_boxes, base_bbox_losses, base_giou_losses, base_iou_scores, 
                                                                                                                                                          pred_boxes, bbox_losses, giou_losses, iou_scores):
            gt_bbox *= torch.tensor([image_size[0], image_size[1], image_size[0], image_size[1]])
            base_bbox *= torch.tensor([image_size[0], image_size[1], image_size[0], image_size[1]])
            pred_bbox *= torch.tensor([image_size[0], image_size[1], image_size[0], image_size[1]])

            visualize_bbox(image_path, gt_bbox, base_bbox, base_bbox_loss, base_giou_loss, base_iou_score, pred_bbox, bbox_loss, giou_loss, iou_score, store_path=store_path)

    return pred_logits, loss, loss_dict


def train_loop(dataset: str,
               model: callable,
               processor: callable,
               data_loader: DataLoader,
               device: str,
               optimizer: torch.optim.Optimizer = None,
               num_classes: int = None,
               train_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
               log_interval: int = 10,
               epoch: int = 0,
               rank: int = 0,
               wandbLogger: wandb.wandb_sdk.wandb_run.Run = None,
               precompute: bool = True,
               eval_only: bool = False,
               weight_dict: dict = None):

    # to pretend training script as evaluation script, need separate evaluation script for faster inference.
    runtime = 'val' if eval_only else 'train'
    if args.eval_test:
        runtime = 'test'

    model.train()
    model.to(device)
    local_model = model.module if world_size > 1 else model

    epoch_loss = 0
    batch_losses = []
    acc_metric_top1 = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes).to(device)
    acc_metric_top5 = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes, top_k=5).to(device)
    accuracy_dict = {f"{runtime}/top1(%)": 100 * 0, f"{runtime}/top5(%)": 100 * 0}
    obj_det_epoch_loss_dict = {}

    # feed parts' names to OwlViT for localization
    with torch.no_grad():
        text_inputs_parts = processor(text=all_parts, padding="max_length", truncation=True, return_tensors="pt").to(device)
        total_descriptors_part = text_inputs_parts['input_ids'].shape[0]
        text_inputs_parts['input_ids'] = text_inputs_parts['input_ids'].repeat(args.batch_size, 1)
        text_inputs_parts['attention_mask'] = text_inputs_parts['attention_mask'].repeat(args.batch_size, 1)

        text_embeds = []
        num_batches = math.ceil(len(all_descriptions) / args.batch_size)
        for i in range(num_batches):
            start = i * args.batch_size
            end = (i+1) * args.batch_size
            text_inputs = processor(text=all_descriptions[start:end], padding="max_length", truncation=True, return_tensors="pt").to(device)
            text_embeds.append(local_model.owlvit.get_text_features(**text_inputs))

        text_embeds_val = []
        num_batches = math.ceil(len(all_descriptions_val) / args.batch_size)
        for i in range(num_batches):
            start = i * args.batch_size
            end = (i+1) * args.batch_size
            text_inputs_val = processor(text=all_descriptions_val[start:end], padding="max_length", truncation=True, return_tensors="pt").to(device)
            text_embeds_val.append(local_model.owlvit.get_text_features(**text_inputs_val))

        text_embeds = torch.cat(text_embeds, dim=0).cpu().detach()
        text_embeds_val = torch.cat(text_embeds_val, dim=0).cpu().detach()

    torch.autograd.set_detect_anomaly(True)

    store_path = None
    visualize_batch_count = 0
    if args.visualize > 0:
        visualize_batch_count = args.visualize

        # Create folder for today's date and time
        store_path = f"../results/bird_soup/visualization/{runtime}_{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}"
        os.makedirs(store_path, exist_ok=True)

    load_boxes_base = args.eval_test and "cub_test_reindexed" in args.test_file
    if load_boxes_base:
        df_cub_test = pd.read_hdf(args.test_file)

    for batch_idx, batch_data in tqdm(enumerate(data_loader), desc=f"{runtime} epoch {epoch}", total=len(data_loader)):
        if not precompute:
            raise NotImplementedError("Current Training script only support precomputed (boxes) mode")

        images, targets_cls, image_paths = batch_data
        images, targets_cls = images.to(device), targets_cls.to(device)
        batch_size = len(image_paths)

        # Handle the last batch separately
        if batch_idx == len(data_loader) - 1 or batch_size != args.batch_size:
            text_inputs_parts['input_ids'] = text_inputs_parts['input_ids'][:total_descriptors_part].repeat(batch_size, 1)
            text_inputs_parts['attention_mask'] = text_inputs_parts['attention_mask'][:total_descriptors_part].repeat(batch_size, 1)

        images['pixel_values'] = images['pixel_values'].squeeze(1).to(device)

        # ------------------------------------------------------------------
        # Increase/Reduce number of classes for contrastive learning
        # ------------------------------------------------------------------
        if args.network_type == "contrastive" and runtime in ["train", "val"]:
            unique_class_ids = list(dict.fromkeys(targets_cls.tolist()).keys())

            # Select X from the remaining classes with X = num_negatives - unique_class_ids
            if runtime == "train":
                if args.contrastive_sampler in ["refilled_empty_classes", "removed_empty_classes"]:
                    assert len(unique_class_ids) == batch_size

                num_extra_negatives = args.num_negatives_train - len(unique_class_ids)
                description_embeds = text_embeds

                if num_extra_negatives > 0:
                    print(f"Adding {num_extra_negatives} extra negatives to the batch for {runtime}")
                    all_negatives = [class_id for class_id in range(len(templated_descriptions.keys())) if class_id not in unique_class_ids]
                    unique_class_ids += random.sample(all_negatives, num_extra_negatives)

            else:
                num_extra_negatives = args.num_negatives_val - len(unique_class_ids)
                description_embeds = text_embeds_val

                if num_extra_negatives > 0:
                    print(f"Adding {num_extra_negatives} extra negatives to the batch for {runtime}")
                    all_negatives = [class_id for class_id in range(len(templated_descriptions_val.keys())) if class_id not in unique_class_ids]
                    unique_class_ids += random.sample(all_negatives, num_extra_negatives)

            selected_text_embeds = description_embeds.view(-1, len(all_parts), description_embeds.shape[-1])[unique_class_ids]
            text_desc_embeds = selected_text_embeds.view(-1, description_embeds.shape[-1]).repeat(batch_size, 1)   #.to(device)

            # Update targets when the order of text_embeds is changed
            targets_cls = torch.tensor([unique_class_ids.index(class_id) for class_id in targets_cls.tolist()]).to(device)

            # Also update number of classes in the model for contrastive loss in the upper branch
            local_model.update_num_classes(len(unique_class_ids))
        else:
            text_desc_embeds = text_embeds.repeat(batch_size, 1)   #.to(device)
        # ------------------------------------------------------------------

        # Update targets for box and class losses in addition to the xclip loss
        image_ids = [".".join(image_path.split("/")[-1].split(".")[:-1]) for image_path in image_paths]
        targets = []
        for idx, image_id in enumerate(image_ids):
            if dataset == "bird_soup":
                boxes_dir = f"{PRECOMPUTED_DIR}/{dataset}/data_updated_v2/{image_id}.pth"
                key_boxes = "boxes_info"    # For precomputed boxes from BirdSoup-v1
                key_logits = "part_logits_owlvit_base"

                boxes_dir_base = None
                if load_boxes_base:
                    ori_image_name = df_cub_test.loc[df_cub_test['new_image_name'] == image_id + ".jpg"]["org_image_name"].values[0]
                    boxes_dir_base = f"{PRECOMPUTED_DIR}/cub/test/owlvit-base-patch32_cub-12-parts/data/{ori_image_name.replace('.jpg', '.pth')}"

            else:
                boxes_dir = f"../pred_boxes/{dataset}/owl_vit_owlvit-large-patch14_descriptors_chatgpt_groundtruths/{image_id}.pth"
                key_boxes = "boxes"
                key_logits = "logits_owlvit_base"

            info_boxes = torch.load(boxes_dir, map_location=device)
            if dataset == "bird_soup" and key_boxes not in info_boxes:
                key_boxes = "part_boxes"    # For precomputed boxes from eBird

            if len(all_parts) != len(info_boxes[key_boxes].keys()):
                info_boxes[key_boxes] = {part: info_boxes[key_boxes][part] for part in all_parts}
                info_boxes[key_logits] = info_boxes[key_logits][:, sel_part_indices]

            assert list(info_boxes[key_boxes].keys()) == all_parts

            # Bounding box coordinates in corner format to [0, 1]
            bboxes = []
            bboxes_ori = []
            image_width, image_height = info_boxes["image_size"]

            for bbox in list(info_boxes[key_boxes].values()):
                bboxes_ori.append(bbox)

                x1, y1, x2, y2 = bbox
                x1_normalized = x1 / image_width
                y1_normalized = y1 / image_height
                x2_normalized = x2 / image_width
                y2_normalized = y2 / image_height
                bboxes.append([x1_normalized, y1_normalized, x2_normalized, y2_normalized])

            bboxes_base = []
            if load_boxes_base:
                info_boxes_base = torch.load(boxes_dir_base, map_location=device)

                for bbox in list(info_boxes_base["boxes_info"].values()):
                    x1, y1, x2, y2 = bbox
                    x1_normalized = x1 / image_width
                    y1_normalized = y1 / image_height
                    x2_normalized = x2 / image_width
                    y2_normalized = y2 / image_height
                    bboxes_base.append([x1_normalized, y1_normalized, x2_normalized, y2_normalized])

            targets.append({"class_labels": torch.tensor(list(range(len(all_parts)))).to(device),
                            "boxes_ori": bboxes_ori,                # 12 boxes from OwlViT-large
                            "boxes": bboxes,                        # 12 boxes from OwlViT-large - normalized
                            "boxes_base": bboxes_base,              # 12 boxes from OwlViT-base - normalized                            
                            "logits": info_boxes[key_logits],       # box logits (576, 12) from OwlViT-base
                            "targets_cls": targets_cls[idx],
                            "image_size": info_boxes["image_size"],
                            "image_path": info_boxes["image_path"],
                            })

        if eval_only:
            with torch.no_grad():
                logits, loss, loss_dict = forward_inputs(local_model, images, text_inputs_parts, text_desc_embeds, targets, weight_dict, visualize_batch_count, store_path)

            batch_loss = loss.item()
            batch_losses.append(batch_loss)
        else:
            optimizer.zero_grad()
            logits, loss, loss_dict = forward_inputs(local_model, images, text_inputs_parts, text_desc_embeds, targets, weight_dict, visualize_batch_count, store_path)

            # compute loss and update model's weights
            batch_loss = loss.item()
            batch_losses.append(batch_loss)

            loss.backward()
            optimizer.step()
            if train_scheduler is not None:
                train_scheduler.step()

        # log model accuracy
        if acc_metric_top1.num_classes == logits.shape[-1]:
            acc_metric_top1.update(logits, targets_cls)
            acc_metric_top5.update(logits, targets_cls)
        else:
            acc_metric_top1.update(logits.argmax(dim=-1), targets_cls)
            # acc_metric_top5.update(torch.tensor([-1] * len(targets_cls)), targets_cls)    # dummy values

        # a dictionary of current accuracy, e.g., {"top1": 0.1, "top5": 0.2}
        accuracy_dict = {f"{runtime}/top1(%)": 100 * acc_metric_top1.compute().item(),
                         f"{runtime}/top5(%)": 100 * acc_metric_top5.compute().item()}

        # od_batch_loss_dict = {f"{runtime}/{k}": v for k, v in loss_dict.items()}
        for key in set(obj_det_epoch_loss_dict.keys()) | set(loss_dict.keys()):
            obj_det_epoch_loss_dict[key] = obj_det_epoch_loss_dict.get(key, []) + [loss_dict[key].item()]

        if visualize_batch_count > 0:
            visualize_batch_count -= 1

        if rank in {-1, 0}:
            if wandbLogger is not None:
                wandbLogger.log(accuracy_dict)
                wandbLogger.log({f"{runtime}/batch_loss": batch_loss})

                # If training to get rid of teacher model => Monitor the loss of the object detection branch
                if not args.logits_from_teacher:
                    wandbLogger.log({f"{runtime}/{key}": values.item() for key, values in loss_dict.items()})

            if batch_idx % log_interval == 0:
                print(f"{runtime} epoch {epoch}: [{batch_idx * len(targets)}/{len(data_loader.dataset)} ({100. * batch_idx / len(data_loader):.0f}%)]\tBatch Loss: {batch_loss:.6f}\tTop1: {accuracy_dict[f'{runtime}/top1(%)']:.2f}")
                print(" ".join([f"{k}: {v:.5f}" for k, v in loss_dict.items()]))

    epoch_loss = sum(batch_losses) / len(batch_losses)
    obj_det_epoch_loss_dict = {f"{runtime}/epoch_{key}": (sum(values) / len(values)) for key, values in obj_det_epoch_loss_dict.items()}

    if eval_only:
        return epoch_loss, accuracy_dict[f'{runtime}/top1(%)'], accuracy_dict[f'{runtime}/top5(%)'], obj_det_epoch_loss_dict

    return model, epoch_loss, accuracy_dict[f'{runtime}/top1(%)'], accuracy_dict[f'{runtime}/top5(%)'], obj_det_epoch_loss_dict


def parse_arguments():
    parser = argparse.ArgumentParser()

    # ------------------------------------------------------------
    #   Must-check arguments for experiments but usually FIXED
    # ------------------------------------------------------------
    parser.add_argument('--model', help='select model', default="owlvit-large-patch14", choices=["owlvit-base-patch32", "owlvit-base-patch16", "owlvit-large-patch14"])
    parser.add_argument('--dataset', help='select dataset', default="cub", choices=["imagenet", "imagenet-v2", "imagenet-a", "imagenet-c", "places365", "cub", "nabirds", "bird_soup"])
    parser.add_argument('--sub_datasets', help='select a group of datasets in Bird Soup', default="all")
    parser.add_argument('--distortion', help='select distortion type if using ImageNet-C', default="defocus_blur", choices=["defocus_blur", "glass_blur", "motion_blur", "zoom_blur", "shot_noise", "gaussian_noise", "impulse_noise"])
    parser.add_argument('--distortion_severity', type=int, help='select distortion severity if using ImageNet-C', default=1, choices=[1, 2, 3, 4, 5])

    parser.add_argument('--epochs', type=int, help='num epochs', default=50)
    parser.add_argument('--batch_size', type=int, help='num training batch size', default=32)
    parser.add_argument('--batch_size_val', type=int, help='num validation batch size for contrastive learning only', default=32)
    parser.add_argument('--eval_size', type=float, help='to select a proportion for val from train set', default=0.2)
    parser.add_argument('--num_workers', type=int, help='num workers for batch processing', default=16)
    parser.add_argument('--num_samples', type=int, help='num images per class', default=-1)
    parser.add_argument('--devices', help='select device', default="0", type=str)
    parser.add_argument('--random_seed', help='random seed (for data subsampling only)', default=42, type=int)
    parser.add_argument('--early_stopping', type=int, help='specify a number greater than 0 if using early stopping', default=0)

    # ------------------------------------------------------------
    #   Must-check arguments for experiments: FREQUENTLY CHANGE
    # ------------------------------------------------------------
    parser.add_argument('--descriptors', help='select descriptors for OwlViT', default="chatgpt", choices=["sachit", "chatgpt"])
    parser.add_argument('--prompt_type', type=int, help='select prompt type', default=5)
    parser.add_argument('--owlvit_threshold', type=float, help='select threshold for owl_vit', default=-1)
    parser.add_argument('--owlvit_conf_scores', help='use owlvit scores as confidence scores', action="store_true")

    parser.add_argument('--lr', type=float, help='learning rate', default=1e-5)
    parser.add_argument('--weight_decay', type=float, help='weight decay', default=0.0001)
    parser.add_argument('--save_freq', type=int, help='save results after a number of epochs', default=0)
    parser.add_argument('--scheduler_mode', type=str, help='select mode for scheduler', default="min")
    parser.add_argument('--scheduler_factor', type=float, help='select factor for scheduler', default=0.5)
    parser.add_argument('--scheduler_patience', type=int, help='select patience for scheduler', default=5)

    parser.add_argument('--num_negatives_train', type=int, help='number of train negatives for contrastive learning', default=32)
    parser.add_argument('--num_negatives_val', type=int, help='number of val negatives for contrastive learning', default=32)

    parser.add_argument('--eval_test', help='run evaluation on test set', action="store_true")
    parser.add_argument('--zeroshot_split', help='split 75% classes for training and 25% classes for testing', action="store_true")
    parser.add_argument("--best_model", type=str, default="", help="load the best model for evaluation.")

    parser.add_argument('--loss_weights', type=str, help='loss weights for loss_ce, loss_bbox, loss_iou, loss_sym_box_label and loss_xclip separated by commas, respectively', default="1,1,1,1,1")
    parser.add_argument('--freeze_box_heads', help='freeze weights of box and class heads of OwlViT', action="store_true")
    parser.add_argument('--train_box_heads_only', help='freeze weights of vision encoder and MLP CLS head of OwlViT', action="store_true")
    parser.add_argument('--network_type', help='select network type for training: contrastive or classification', default=None, choices=["contrastive", "classification"])
    parser.add_argument('--classification_loss', help='select the main loss for image classification', default=None, choices=["ce_loss", "focal_loss"])
    parser.add_argument('--contrastive_sampler', help='select the sampling strategy for contrastive', default=None, choices=["refilled_empty_classes", "removed_empty_classes"])
    parser.add_argument('--logits_from_teacher', help='use boxes logits from teacher model (e.g., OwlViT-base) for 12 box selection ', action="store_true")

    parser.add_argument("--part_names", type=str, default=None, help="select part names for training and evaluation in "
                                                                     "'back,beak,belly,breast,crown,forehead,eyes,legs,wings,nape,tail,throat'. Default to None")

    parser.add_argument('--train_file', type=str, help='set path to the train file', default=BIRD_SOUP_META_PATH_TRAIN)
    parser.add_argument('--val_file', type=str, help='set path to the validation file', default=BIRD_SOUP_META_PATH_VAL)
    parser.add_argument('--test_file', type=str, help='set path to the test file', default=BIRD_SOUP_META_PATH_TEST)
    parser.add_argument('--descriptor_path', type=str, help='if set, descriptors are loaded from file', default=None)

    parser.add_argument('--birdsoup_level', type=int, help='specify bird soup level from 1 to 3 for logging purposes', choices=[1, 2, 3])
    parser.add_argument('--image_mean_std', help='use image mean and std precomputed from CLIP or BirdSoup images', default="bird_soup", choices=["bird_soup", "clip"])
    parser.add_argument('--finetuning', help='finetune pretrained models on downstream tasks', default=None, choices=["vision_encoder_mlp", "mlp_only"])    # "proj_layer", "linear"

    # For FOCAL loss
    parser.add_argument('--alpha', type=float, help='alpha for focal loss', default=0.25)
    parser.add_argument('--gamma', type=float, help='gamma for focal loss', default=2.0)

    parser.add_argument('--box_head_num_layers', type=int, help='define number of linear layers for box head', default=3)
    parser.add_argument('--ablation', help='run ablation study on number of classes for contrastive training', action="store_true")

    # ------------------------------------------------------------
    #   Analysis
    # ------------------------------------------------------------
    parser.add_argument('--visualize', type=int, help='visualization', default=0)
    parser.add_argument('--test_class_list', help='a file containing a list of class to test', default=[], nargs='+')

    # ------------------------------------------------------------
    #   Others
    # ------------------------------------------------------------
    parser.add_argument('--note', type=str, help='write whatever to quickly remind experiment settings', default="")
    parser.add_argument('--debugging', help='do not save results and logs when debugging', action="store_true")
    parser.add_argument('--overwrite_boxes', help='save boxes for tuning', action="store_true")
    parser.add_argument('--check_box_files', help='check if box file exists (otherwise box computing will skip if the box folder exist.)', action="store_true")
    parser.add_argument('--verbose', help='print logs', action="store_true")
    parser.add_argument("--no_log", action='store_true', help="disable wandb logging.")
    parser.add_argument("--project_name", type=str, default="xclip", help="name of the wandb project.")

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_arguments()
    start_time = datetime.now()

    seed_everything(args.random_seed)

    # load general configs for training
    world_size = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    global_rank = int(os.environ["RANK"]) if "RANK" in os.environ else -1 # we only have one node so the global rank should always be 0

    if global_rank != -1:
        rank = int(os.environ["LOCAL_RANK"])
        dist.init_process_group("nccl", init_method="env://")
        mp.set_start_method("spawn", force=True)  # change to spawn to avoid issues with pytorch dataloader when num_workers > 0
    else:
        rank = -1

    wandbLogger = None
    if not args.no_log or args.eval_test:
        run_name = f"{str(datetime.now().strftime('%m_%d_%Y-%H:%M:%S'))}_lr_{args.lr}_{args.epochs}ep_prompt{args.prompt_type}"
        sub_dir = f"evaluation_{args.network_type}" if args.eval_test else f"training_{args.network_type}"
        if args.finetuning is not None:
            sub_dir = f"finetune_{args.network_type}/{args.finetuning}"

        out_dir = f'{PROJECT_ROOT}/results/{args.dataset}/level_{args.birdsoup_level}/{sub_dir}/{args.descriptors}-{args.model.replace("/", "")}/{run_name}'
        if args.note != "":
            out_dir += f'_{args.note}'
            run_name += f'_{args.note}'

        if rank in {-1, 0}:
            os.makedirs(out_dir, exist_ok=True)

            if not args.no_log:
                wandbLogger = wandb.init(project=args.project_name, name=run_name, resume=False, dir=out_dir, mode='disabled' if args.no_log else 'online', entity='pmthangk09')
            else:
                print("Warning: wandb logging is disabled. Make sure this action is intended.")

            # Save the modified version of the config files
            with open(f'{out_dir}/configs.json', 'w') as f:
                json.dump(args.__dict__, f, indent=4)

    # split the devices into a list and convert to ints
    device_list = [int(x) for x in args.devices.split(",")]
    check_device_availability(device_list)
    device = f'cuda:{device_list[rank]}' if world_size > 1 else f'cuda:{device_list[0]}'
    torch.cuda.set_device(device)

    # load pre-trained model
    owlvit_det_processor = OwlViTProcessor.from_pretrained(f"google/{args.model}")
    owlvit_det_model = OwlViTForObjectDetection.from_pretrained(f"google/{args.model}")

    # TODO: Update MEAN + STD for image normalization
    ''' CLIP VALUES
    image_mean = {list} [0.48145466, 0.4578275, 0.40821073]
    image_std = {list} [0.26862954, 0.26130258, 0.27577711]
    '''
    if args.image_mean_std == "bird_soup":
        owlvit_det_processor.image_processor.image_mean = [0.48168647, 0.49244233, 0.42851609]
        owlvit_det_processor.image_processor.image_std = [0.18656386, 0.18614962, 0.19659419]

    # load dataset
    # The order/indices of classes in target_classes MUST BE the same as the order/indices of classes in the dataset
    train_dataset, val_dataset = load_training_dataset(args.dataset, args.sub_datasets, args.eval_size, transform=owlvit_det_processor, random_state=args.random_seed, zeroshot_split=args.zeroshot_split)
    target_classes = train_dataset.classes if hasattr(train_dataset, "classes") else train_dataset.dataset.classes
    target_classes_val = val_dataset.classes if hasattr(val_dataset, "classes") else val_dataset.dataset.classes

    test_dataset = None
    if args.eval_test:
        test_dataset, _ = load_training_dataset(args.dataset, args.sub_datasets, 1.0, transform=owlvit_det_processor, random_state=args.random_seed, zeroshot_split=args.zeroshot_split, split="test")
        target_classes = test_dataset.classes if hasattr(test_dataset, "classes") else test_dataset.dataset.classes

    if args.dataset == "bird_soup":
        target_classes = [c.lower().replace("-", " ").replace("'s", "") for c in target_classes]
        target_classes_val = [c.lower().replace("-", " ").replace("'s", "") for c in target_classes_val]

    # prepare text embeddings
    # Use target_classes to filter out classes that are not in the dataset (for BirdSoup)
    descriptions_only, _ = load_descriptions(dataset_name=args.dataset, prompt_type=0, desc_type=args.descriptors, target_classes=target_classes, descriptor_path=args.descriptor_path)
    templated_descriptions, _ = load_descriptions(dataset_name=args.dataset, prompt_type=args.prompt_type, desc_type=args.descriptors, target_classes=target_classes, descriptor_path=args.descriptor_path)
    templated_descriptions_val, _ = load_descriptions(dataset_name=args.dataset, prompt_type=args.prompt_type, desc_type=args.descriptors, target_classes=target_classes_val, descriptor_path=args.descriptor_path)

    # Sorted the keys in templated_descriptions to match the order of classes in target_classes
    templated_descriptions = {k: templated_descriptions[k] for k in sorted(templated_descriptions, key=target_classes.index)}
    templated_descriptions_val = {k: templated_descriptions_val[k] for k in sorted(templated_descriptions_val, key=target_classes_val.index)}

    # The class names in class_list must match target_classes to correctly map the descriptors to the classes
    assert list(templated_descriptions.keys()) == target_classes
    assert list(templated_descriptions_val.keys()) == target_classes_val

    num_classes = len(templated_descriptions.keys())
    all_descriptions = list(templated_descriptions.values())
    all_descriptions_val = list(templated_descriptions_val.values())

    # Use parts only for localization
    all_parts = []
    if args.descriptors == "chatgpt":
        all_parts = [[descriptor.split(":")[0] for descriptor in descriptors if ":" in descriptor] for descriptors in descriptions_only.values()][0]
        sel_part_indices = list(range(len(all_parts)))

        # The order of part names is still preserved
        if args.part_names is not None:
            part_names = args.part_names.split(",")
            sel_part_indices = [i for i, part in enumerate(all_parts) if part in part_names]
            all_parts = [part for part in all_parts if part in part_names]

        all_descriptions = [[descriptor for i, descriptor in enumerate(descriptors) if i in sel_part_indices] for descriptors in all_descriptions]
        all_descriptions_val = [[descriptor for i, descriptor in enumerate(descriptors) if i in sel_part_indices] for descriptors in all_descriptions_val]

    # TODO: For contrastive training, target classes of train and validation are the same EXCEPT FOR ablation study
    # Ablation study uses a separate validation set whose number of classes is different from the train set
    # The number of classes in the validation set for normal training is slightly less than the train set
    # because some classes in the train set do not have enough samples for validation (i.e., >= 3 images)
    if not args.ablation:
        all_descriptions_val = all_descriptions

    # Loss weights for training
    loss_weights = [float(weight) for weight in args.loss_weights.split(",")]
    weight_dict = {"loss_ce": loss_weights[0], "loss_bbox": loss_weights[1], "loss_giou": loss_weights[2],
                   "loss_sym_box_label": loss_weights[3], "loss_xclip": loss_weights[4]}

    # Initialize OwlViT model for Classification
    model = OwlViTForClassification(owlvit_det_model=owlvit_det_model, num_classes=num_classes, num_parts=len(all_parts), device=device,
                                    freeze_box_heads=args.freeze_box_heads, train_box_heads_only=args.train_box_heads_only,
                                    network_type=args.network_type, classification_loss=args.classification_loss,
                                    weight_dict=weight_dict, logits_from_teacher=args.logits_from_teacher,
                                    finetuning=args.finetuning, alpha=args.alpha, gamma=args.gamma)

    trained_params, frozen_params = 0, 0
    for name, param in model.named_parameters():
        if not args.no_log and rank in {-1, 0}:
            print(f"{name}: {param.shape if len(param.shape) > 0 else param.type()} - Required grad: {param.requires_grad}")

        params = reduce(operator.mul, list(param.shape)) if len(param.shape) > 0 else 1
        if param.requires_grad:
            trained_params += params
        else:
            frozen_params += params

    print(f"Trainable parameters: {format(trained_params, ',')}")
    print(f"Frozen parameters: {format(frozen_params, ',')}")

    if args.best_model:
        # Load best model: OwlViT-Base + OwlViT_CLS_PLUS + Finetune MLP + Vision encoder
        ckpt = torch.load(args.best_model, map_location=device)
        model.load_state_dict(ckpt, strict=False)

    # TODO: FORCE UPDATING BOX HEAD
    if args.box_head_num_layers > 3:
        model.update_box_head(args.box_head_num_layers)

    # Write model architecture to file
    if not args.no_log:
        # write model architecture to file for tracking purposes
        with open(f'{out_dir}/model_arch.txt', 'w') as f:
            f.write(str(model) + "\n\n")

            for name, param in model.named_parameters():
                f.write(f"{name}: {param.shape if len(param.shape) > 0 else param.type()} - Required grad: {param.requires_grad}\n")

            f.write(f"\nTrainable parameters: {format(trained_params, ',')}\n")
            f.write(f"Frozen parameters: {format(frozen_params, ',')}\n")

    if world_size > 1:
        # TODO: Batch_size 200 (i.e., 200 images per batch) OR Batch_size 32 + num_negatives = 200 ???
        train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
        val_sampler = DistributedSampler(val_dataset, num_replicas=world_size, rank=rank, shuffle=True)
        test_sampler = DistributedSampler(test_dataset, num_replicas=world_size, rank=rank, shuffle=False) if test_dataset else None
        train_shuffle = val_shuffle = False
    else:
        train_sampler, val_sampler, test_sampler = None, None, None
        train_shuffle = val_shuffle = True

    test_shuffle = True
    # test_shuffle = False
    # if args.network_type == "contrastive":
    #     test_shuffle = True

    # prepare dataloaders
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers, sampler=train_sampler, shuffle=train_shuffle, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size_val, num_workers=args.num_workers, sampler=val_sampler, shuffle=val_shuffle, pin_memory=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=args.num_workers, sampler=test_sampler, shuffle=test_shuffle, pin_memory=True) if args.eval_test else None

    # SyncBatchNorm
    if world_size > 1 and rank != -1:
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

        # TODO: Check if the data is randomly distributed across the GPUs
        # for i, (input, target) in enumerate(train_loader):
        #     # If the GPU id is accessible as variable `gpu`
        #     if i < 10:  # print only for first 10 batches
        #         print(f"GPU: {rank}, Batch: {i}, First item ID: {input[0]['id']}")

    if rank != -1:
        model = DDP(model, device_ids=[device], output_device=device)
    else:
        model.to(device)

    # if resume:
        # load the latest checkpoint
        # resume_path = train_config["log_dir"] if resume_path is None else resume_path
        # checkpoint = torch.load(os.path.join(resume_path, "last.pt"))
        # model.load_state_dict(checkpoint["model"])

    # load optimizer
    optimizer = torch.optim.AdamW(model.parameters(), lr=float(args.lr), weight_decay=float(args.weight_decay))

    # use plateau scheduler
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=args.scheduler_mode, factor=args.scheduler_factor, patience=args.scheduler_patience, verbose=args.verbose)

    # train loops
    if not args.eval_test:
        val_best_acc, val_best_loss, best_epoch = 0, 9999, 0
        early_stopping = 0

        for epoch in range(args.epochs):
            # train
            train_results = train_loop(dataset=args.dataset, model=model, processor=owlvit_det_processor, data_loader=train_loader,
                                       num_classes=num_classes, device=device, rank=rank, wandbLogger=wandbLogger,
                                       optimizer=optimizer, log_interval=10, epoch=epoch, weight_dict=weight_dict)

            model, train_loss, train_top1, train_top5, train_od_epoch_loss_dict = train_results

            # eval
            eval_results = train_loop(dataset=args.dataset, model=model, processor=owlvit_det_processor, data_loader=val_loader,
                                      num_classes=num_classes, device=device, rank=rank,  wandbLogger=wandbLogger,
                                      eval_only=True, weight_dict=weight_dict)

            val_loss, val_top1, val_top5, val_od_epoch_loss_dict = eval_results

            if world_size > 1:
                # reduce the validation loss across all processes
                dist.all_reduce(torch.tensor(val_loss).to(device), op=dist.ReduceOp.SUM)
            scheduler.step(val_loss)

            if rank in {-1, 0}:
                train_log_dict = {"train/epoch_loss": train_loss, 'train/epoch_top1': train_top1, 'train/epoch_top5': train_top5}
                val_log_dict = {"val/epoch_loss": val_loss, 'val/epoch_top1': val_top1, 'val/epoch_top5': val_top5}
                
                if not args.no_log:
                    wandb.log(train_log_dict | val_log_dict | train_od_epoch_loss_dict | val_od_epoch_loss_dict | {"epochs": epoch}, commit=False)

                sd = model.module.state_dict() if world_size > 1 else model.state_dict()
                if args.save_freq > 0 and epoch % args.save_freq == 0:
                    torch.save(sd, os.path.join(wandb.run.dir, f"e{epoch:04d}.pt"))

                if val_loss < val_best_loss:
                    val_best_loss = val_loss
                    val_best_acc = val_top1
                    best_epoch = epoch
                    torch.save(sd, os.path.join(wandb.run.dir, "best.pt"))
                    early_stopping = 0
                else:
                    early_stopping += 1

                torch.save(sd, os.path.join(wandb.run.dir, "last.pt"))

            if 0 < args.early_stopping <= early_stopping:
                print("EARLY STOPPING TRIGGERED... QUIT TRAINING")
                break

        best_acc = val_best_acc
        best_loss = val_best_loss

        print(f'Best val/val_loss top1: {val_best_acc:.4f}/{val_best_loss:.4f} at epoch {best_epoch}.')
        print(" ".join([f"{k}: {v:.5f}" for k, v in val_od_epoch_loss_dict.items()]))
    else:
        test_results = train_loop(dataset=args.dataset, model=model, processor=owlvit_det_processor, data_loader=test_loader,
                                  num_classes=num_classes, device=device, rank=rank, wandbLogger=wandbLogger,
                                  eval_only=True, weight_dict=weight_dict)

        test_loss, test_top1, test_top5, test_od_epoch_loss_dict = test_results

        best_acc = test_top1
        best_loss = test_loss
        best_epoch = None

        print(f'Best test / loss top1: {best_acc:.4f} / {best_loss:.4f}.')
        print(" ".join([f"{k}: {v:.5f}" for k, v in test_od_epoch_loss_dict.items()]))

    end_time = datetime.now()

    with open(f'{out_dir}/results.json', 'w') as f:
        json.dump({f"Best Accuracy ({'val' if not args.eval_test else 'test'})": best_acc,
                   "Best epoch": best_epoch,
                   "Number of examples": len(val_dataset),
                   "Start time": start_time.strftime("%d/%m/%Y %H:%M:%S"),
                   "End time": end_time.strftime("%d/%m/%Y %H:%M:%S"),
                   "Duration": round((end_time - start_time).total_seconds() * 1.0 / 3600, 2)}, f, indent=4)


src/train_owlvit.py
anguyen8/peeb
import argparse
import os
from datetime import datetime

from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import OwlViTProcessor, OwlViTForObjectDetection

from data_loader import DatasetWrapper
from utils import *


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # ------------------------------------------------------------
    #   Must-check arguments for experiments but usually FIXED
    # ------------------------------------------------------------
    parser.add_argument('--model', help='select model', default="owlvit-large-patch14", choices=["owlvit-base-patch32", "owlvit-base-patch16", "owlvit-large-patch14"])
    parser.add_argument('--dataset', help='select dataset', default="imagenet", choices=["imagenet", "imagenet-v2", "imagenet-a", "imagenet-c", "places365", "cub", "nabirds", "bird_soup"])
    parser.add_argument('--distortion', help='select distortion type if using ImageNet-C', default="defocus_blur", choices=["defocus_blur", "glass_blur", "motion_blur", "zoom_blur", "shot_noise", "gaussian_noise", "impulse_noise"])
    parser.add_argument('--distortion_severity', type=int, help='select distortion severity if using ImageNet-C', default=1, choices=[1, 2, 3, 4, 5])

    parser.add_argument('--batch_size', type=int, help='num batch size', default=32)
    parser.add_argument('--num_workers', type=int, help='num workers for batch processing', default=16)
    parser.add_argument('--num_samples', type=int, help='num images per class', default=-1)
    parser.add_argument('--device', help='select device', default="cuda:0", type=str)
    parser.add_argument('--random_seed', help='random seed (for data subsampling only)', default=42, type=int)
    parser.add_argument('--image_type_cub', help='select image type for CUB only', type=str, default=None, choices=["synthetic", "upsampler"])
    parser.add_argument('--box_type_cub', help='select pred_box type for CUB only', type=str, default=None, choices=["with_bird_box", "synthetic", "upsampler"])

    parser.add_argument('--descriptors', help='select descriptors for OwlViT', default="sachit", choices=["sachit", "chatgpt"])
    parser.add_argument('--prompt_type', type=int, help='select prompt type', default=5)
    parser.add_argument('--owlvit_threshold', type=float, help='select threshold for owl_vit', default=-1)

    args = parser.parse_args()

    start_time = datetime.now()

    if not torch.cuda.is_available():
        device = 'cpu'
    else:
        check_device_availability(args.device)
        device = args.device

    boxes_dir = f"{PRECOMPUTED_DIR}/{args.dataset}/data"

    updated_boxes_dir = boxes_dir + "_updated_v2"
    if not os.path.exists(updated_boxes_dir):
        os.makedirs(updated_boxes_dir)

    if args.dataset == "cub":
        boxes_dir += "" if args.box_type_cub is None else f"_{args.box_type_cub}"
    elif args.dataset == "imagenet-c":
        boxes_dir = f"{boxes_dir}/distortion_{args.distortion}_severity_level_{args.distortion_severity}"

    # Prepare text embeddings
    descriptions_only, _ = load_descriptions(dataset_name=args.dataset, prompt_type=0, desc_type=args.descriptors)

    # Remove 'It has' and 'It is' from Sachit's CUB descriptors
    if args.dataset == "cub":
        descriptions_only = {key: [value.replace("It has", "").replace("It is", "").strip() for value in values] for key, values in descriptions_only.items()}

    class_list = list(descriptions_only.keys())
    all_descriptions = list(descriptions_only.values())
    if args.descriptors == "chatgpt":
        all_descriptions = [[descriptor.split(":")[0] for descriptor in descriptors if ":" in descriptor] for descriptors in descriptions_only.values()][0]

    processor = OwlViTProcessor.from_pretrained(f"google/owlvit-base-patch32")
    model = OwlViTForObjectDetection.from_pretrained(f"google/owlvit-base-patch32")
    model.to(device)
    model.eval()

    # Load dataset and prepare data loader
    dataset_args = {"image_type_cub": args.image_type_cub}
    dataset = DatasetWrapper(dataset_name=args.dataset, transform=processor,
                             distortion=f"{args.distortion}:{args.distortion_severity}",
                             samples_per_class=args.num_samples, random_seed=args.random_seed, **dataset_args)

    dataloader = DataLoader(dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)

    with torch.no_grad():
        text_inputs = processor(text=all_descriptions, padding="max_length", truncation=True, return_tensors="pt").to(device)
        total_descriptors = text_inputs['input_ids'].shape[0]
        text_inputs['input_ids'] = text_inputs['input_ids'].repeat(args.batch_size, 1)
        text_inputs['attention_mask'] = text_inputs['attention_mask'].repeat(args.batch_size, 1)

    for batch_idx, batch in tqdm(enumerate(dataloader), desc='Localizing descriptors', total=len(dataloader)):
        images, gt_labels, image_paths, image_ids, image_sizes = batch

        # Handle the last batch separately
        if batch_idx == len(dataloader) - 1:
            text_inputs['input_ids'] = text_inputs['input_ids'][:len(image_paths) * total_descriptors]
            text_inputs['attention_mask'] = text_inputs['attention_mask'][:len(image_paths) * total_descriptors]

        images['pixel_values'] = images['pixel_values'].squeeze(1).to(device)
        with torch.no_grad():
            owl_inputs = images | text_inputs
            owl_outputs = model(**owl_inputs)

            logits = torch.sigmoid(owl_outputs.logits).detach().cpu()

            for i, logit_scores in enumerate(logits):
                results = torch.load(f"{boxes_dir}/{image_ids[i]}.pth")

                # Do not overwrite if already exists
                if 'part_logits_owlvit_base' in results:
                    continue

                results['part_logits_owlvit_base'] = logit_scores
                torch.save(results, f"{updated_boxes_dir}/{image_ids[i]}.pth")




src/update_boxes_logits.py
anguyen8/peeb
import os
import clip
import json
import random
import shutil

from PIL import Image
from torch.nn import functional as F
import torchvision.transforms as transforms

from configs import *


''' CLIP TRANSFORM
0 = {Resize} Resize(size=224, interpolation=bicubic, max_size=None, antialias=None)
1 = {CenterCrop} CenterCrop(size=(224, 224))
2 = {function} <function _convert_image_to_rgb at 0x7fd16ab12830>
3 = {ToTensor} ToTensor()
4 = {Normalize} Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
'''

def img_transform(n_px):
    return transforms.Compose([
        transforms.Resize(n_px, interpolation=Image.BICUBIC),
        transforms.CenterCrop(n_px),
        lambda image: image.convert("RGB"),
        transforms.ToTensor(),
        transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])


def load_descriptions(dataset_name, classes_to_load=None, prompt_type=None, desc_type="sachit", part_based=False,
                      target_classes: list[int] = None, descriptor_path: str = None):
    templated_descriptions, descriptions_mappings = {}, {}

    # ImageNet and ImageNet-v2 share the same list of descriptions
    dataset_to_load = "imagenet" if dataset_name in ["imagenet-v2", "imagenet-a", "imagenet-c"] else dataset_name

    if descriptor_path is None:
        descriptor_path = f"{PROJECT_ROOT}/data/descriptors/{desc_type}/descriptors_{dataset_to_load}.json"

        if dataset_to_load == "bird_11K":
            descriptor_path = descriptor_path.replace(dataset_to_load, f"{dataset_to_load}")

    print("Using descriptors from: ", descriptor_path)

    with open(descriptor_path) as input_file:
        descriptions = json.load(input_file)

    if classes_to_load is not None:
        descriptions = {c: descriptions[c] for c in classes_to_load}
    elif dataset_name == "imagenet-a":
        descriptions = {c: descriptions[c] for idx, c in enumerate(descriptions.keys()) if idx in indices_in_1k}

    if prompt_type is not None:
        for idx, (class_name, class_descriptors) in enumerate(descriptions.items()):
            if target_classes is not None and \
               (class_name.lower() not in target_classes or class_name.lower().replace("-", " ") not in target_classes):
                continue

            if len(class_descriptors) == 0:
                class_descriptors = ['']

            class_name = wordify(class_name)

            # Sachit's prompt
            if prompt_type == 0:
                templated_descriptors = class_descriptors
            elif prompt_type == 1:
                templated_descriptors = [f"{make_descriptor_sentence(class_name, descriptor, part_based)}" for descriptor in class_descriptors]
            elif prompt_type == 2:
                templated_descriptors = [f"{descriptor} of {class_name}" for descriptor in class_descriptors]
            elif prompt_type == 3:
                templated_descriptors = [f"a photo of {descriptor} of {class_name}" for descriptor in class_descriptors]
            elif prompt_type == 4:
                templated_descriptors = [f"a cropped photo of {descriptor} of {class_name}" for descriptor in class_descriptors]
            elif prompt_type == 5:
                templated_descriptors = [f"a photo of a {make_descriptor_sentence(class_name, descriptor, part_based)}" for descriptor in class_descriptors]
            elif prompt_type == 6:
                templated_descriptors = [f"a cropped photo of a {make_descriptor_sentence(class_name, descriptor, part_based)}" for descriptor in class_descriptors]
            elif prompt_type == 7:
                templated_descriptors = [f"a photo of a {make_descriptor_sentence('', descriptor, part_based)}" for descriptor in class_descriptors]
            # Used for finetuning CLIP
            elif prompt_type == 8:
                template = "a {} {} of {}."
                templated_descriptors = [template.format(descriptor.split(":")[1].strip(), descriptor.split(":")[0].strip(), class_name) for descriptor in class_descriptors]
            elif prompt_type == 9:      # Same as Peijie's prompt X-2
                template = "a {} {}"    # convert '{part}: {features}' to 'a {features} {part}'
                templated_descriptors = [template.format(descriptor.split(":")[1].strip(), descriptor.split(":")[0].strip()) for descriptor in class_descriptors]
            elif prompt_type == 10:     # Same as Peijie's prompt X-3
                template = "{}. {}."    # convert '{part}: {features}' to '{features}. {part}.'
                templated_descriptors = [template.format(descriptor.split(":")[1].strip(), descriptor.split(":")[0].strip()) for descriptor in class_descriptors]
            else:
                raise ValueError("Unknown prompt type")

            templated_descriptions[class_name] = templated_descriptors
            descriptions_mappings[class_name] = {templated_descriptor: descriptor for descriptor, templated_descriptor in zip(class_descriptors, templated_descriptors)}

            # Print an example for checking
            if idx == 0:
                print(f"\nExample description for prompt type {prompt_type}: \"{templated_descriptions[class_name][0]}\"\n")

    return templated_descriptions, descriptions_mappings


def wordify(string):
    word = string.replace('_', ' ')
    return word


def make_descriptor_sentence(class_name, descriptor, part_based=False):

    if part_based:
        part, descriptor = descriptor.split(":")
        descriptor = descriptor.strip()

        if class_name != "":
            if descriptor.startswith(('a', 'an', 'used')):
                return f"{class_name}, whose {part} is {descriptor}"
            elif descriptor.startswith(('has', 'often', 'typically', 'may', 'can')):
                return f"{class_name}, whose {part} {descriptor}"
            else:
                return f"{class_name}, whose {part} has {descriptor}"
        else:
            if part == "bird":
                return descriptor

            return f"{part}, which is {descriptor}"

    if descriptor.startswith(('a', 'an', 'used')):
        return f"{class_name}, which is {descriptor}"
    elif descriptor.startswith(('has', 'often', 'typically', 'may', 'can')):
        return f"{class_name}, which {descriptor}"
    else:
        return f"{class_name}, which has {descriptor}"


def check_device_availability(devices: list[int or str] or int or str):
    # check if the devices are in format 'cuda:X' or 'X' or X, where X is an integer
    def _check_device_format(device_name: str or int) -> int:
        if isinstance(device_name, int):
            return device_name
        elif isinstance(device_name, str):
            if device_name.startswith("cuda:"):
                return int(device_name.split(":")[1])
            else:
                return int(device_name)      
    
    # check devices if devices is a list
    if not isinstance(devices, list):
        devices = [devices]
    devices = [_check_device_format(device) for device in devices]
    # check if the devices are available
    
    available_list = list(range(torch.cuda.device_count()))
    # check if the device is available
    for device in devices:
        if device not in available_list:
            raise ValueError(f"Device {device} is not available. Available devices are {available_list}")




src/utils.py
anguyen8/peeb