provided script to prepare cub & pet dataset

.gitignore

@@ -7,4 +7,5 @@ pretrained_weights/
 data/images/
 visualization/
 data/annotations/
-=*
+=*
+*.tgz

README.md

@@ -77,21 +77,44 @@ cub/
 
 ## CUB
 
-- Download prepared dataset
-    - From [![](https://img.shields.io/badge/google_drive-yellow)](https://drive.google.com/drive/folders/1X3ikQEk_D7cKcyCnxbF3kJTsZ0LZfvVO?usp=sharing)
-- `Or` Prepare the dataset by yourself
-    - You can download the CUB dataset from [the original website](https://www.vision.caltech.edu/datasets/cub_200_2011/) and put it in the `data/images/` folder.
-    - You can use the dataset's provided train/val split to create the train/val splits and have their class numbers as the `prefix` of the respective image folder names(starting from 1).
-    - The code will automatically create train and val annotation files in the `data/annotations/` folder for each dataset if not provided.
+1. Download `CUB_200_2011.tgz` from [the official website](https://www.vision.caltech.edu/datasets/cub_200_2011/) and extract it:
+    ```bash
+    tar -xzf CUB_200_2011.tgz
+    ```
+    You will get a `CUB_200_2011/` folder containing `images/`, `images.txt`, and `train_test_split.txt`.
+
+2. Run the preparation script to arrange the images into the required structure:
+    ```bash
+    python data/prepare_cub.py \
+        --cub_dir /path/to/CUB_200_2011 \
+        --out_dir /path/to/data/images/cub
+    ```
+
+The script uses the official train/test split and organises images into `cub/train/` and `cub/val/` with class folders named `NNN.ClassName` (e.g. `001.Black_footed_Albatross`).
 
 </details>
 <details>
 <summary>Prepare Oxford Pet dataset</summary>
 
 ## Pet Dataset
-- Download prepared dataset
-    - From [![](https://img.shields.io/badge/google_drive-yellow
-)](https://drive.google.com/drive/folders/1X3ikQEk_D7cKcyCnxbF3kJTsZ0LZfvVO?usp=sharing)
+
+1. Download both archives from [the official website](https://www.robots.ox.ac.uk/~vgg/data/pets/) and extract them:
+    ```bash
+    tar -xzf images.tar.gz
+    tar -xzf annotations.tar.gz
+    ```
+    You will get an `images/` folder with all `.jpg` files and an `annotations/` folder containing `trainval.txt` and `test.txt`.
+
+2. Run the preparation script to arrange the images into the required structure:
+    ```bash
+    python data/prepare_pet.py \
+        --images_dir /path/to/images \
+        --annotations_dir /path/to/annotations \
+        --out_dir /path/to/data/images/pet
+    ```
+
+The script maps `trainval.txt` → `pet/train/` and `test.txt` → `pet/val/`, with class folders named `NNN.BreedName` (e.g. `001.Abyssinian`).
+
 </details>
 
 **To add new dataset, see [Extensions](#extensions)**

data/prepare_cub.py

@@ -0,0 +1,90 @@
+"""
+Prepare CUB-200-2011 dataset into the Prompt-CAM directory structure.
+
+Download the dataset from:
+    https://www.vision.caltech.edu/datasets/cub_200_2011/
+
+After extracting the tarball you should have:
+    CUB_200_2011/
+    ├── images/
+    ├── images.txt
+    └── train_test_split.txt
+
+Run:
+    python data/prepare_cub.py --cub_dir /path/to/CUB_200_2011 --out_dir /path/to/data/images/cub
+"""
+
+import argparse
+import os
+import shutil
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cub_dir",
+        required=True,
+        help="Path to the extracted CUB_200_2011 folder.",
+    )
+    parser.add_argument(
+        "--out_dir",
+        required=True,
+        help="Destination folder (e.g. data/images/cub). Created if absent.",
+    )
+    return parser.parse_args()
+
+
+def read_split(cub_dir):
+    """Return dict {image_id: 'train' or 'val'}."""
+    split_file = os.path.join(cub_dir, "train_test_split.txt")
+    split = {}
+    with open(split_file) as f:
+        for line in f:
+            img_id, is_train = line.strip().split()
+            split[img_id] = "train" if is_train == "1" else "val"
+    return split
+
+
+def read_images(cub_dir):
+    """Return dict {image_id: relative_path_from_images_dir}."""
+    images_file = os.path.join(cub_dir, "images.txt")
+    images = {}
+    with open(images_file) as f:
+        for line in f:
+            img_id, img_path = line.strip().split(maxsplit=1)
+            images[img_id] = img_path
+    return images
+
+
+def main():
+    args = parse_args()
+    cub_dir = args.cub_dir
+    out_dir = args.out_dir
+
+    split = read_split(cub_dir)
+    images = read_images(cub_dir)
+
+    images_src = os.path.join(cub_dir, "images")
+
+    for img_id, img_rel_path in images.items():
+        subset = split[img_id]
+        # img_rel_path is like "001.Black_footed_Albatross/Black_Footed_Albatross_0001_796111.jpg"
+        class_folder = img_rel_path.split("/")[0]
+        filename = img_rel_path.split("/")[1]
+
+        dst_dir = os.path.join(out_dir, subset, class_folder)
+        os.makedirs(dst_dir, exist_ok=True)
+
+        src = os.path.join(images_src, img_rel_path)
+        dst = os.path.join(dst_dir, filename)
+        if not os.path.exists(dst):
+            shutil.copy2(src, dst)
+
+    train_classes = len(os.listdir(os.path.join(out_dir, "train")))
+    val_classes = len(os.listdir(os.path.join(out_dir, "val")))
+    print(f"Done. train: {train_classes} classes | val: {val_classes} classes")
+    print(f"Output saved to: {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

data/prepare_pet.py

@@ -0,0 +1,122 @@
+"""
+Prepare Oxford-IIIT Pet dataset into the Prompt-CAM directory structure.
+
+Download the dataset from:
+    https://www.robots.ox.ac.uk/~vgg/data/pets/
+
+You need two archives:
+    images.tar.gz        -> extract to get a folder of .jpg images
+    annotations.tar.gz   -> extract to get trainval.txt and test.txt
+
+After extracting you should have:
+    <some_dir>/
+    ├── images/               (all .jpg images, e.g. Abyssinian_1.jpg)
+    └── annotations/
+        ├── trainval.txt
+        └── test.txt
+
+Each annotation line has the format:
+    <image_name> <class_id> <species_id> <breed_id>
+where class_id is 1-indexed and corresponds to alphabetical order of breed names.
+
+Run:
+    python data/prepare_pet.py \
+        --images_dir /path/to/images \
+        --annotations_dir /path/to/annotations \
+        --out_dir /path/to/data/images/pet
+"""
+
+import argparse
+import os
+import re
+import shutil
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--images_dir",
+        required=True,
+        help="Path to the extracted images/ folder containing all .jpg files.",
+    )
+    parser.add_argument(
+        "--annotations_dir",
+        required=True,
+        help="Path to the extracted annotations/ folder containing trainval.txt and test.txt.",
+    )
+    parser.add_argument(
+        "--out_dir",
+        required=True,
+        help="Destination folder (e.g. data/images/pet). Created if absent.",
+    )
+    return parser.parse_args()
+
+
+def class_name_from_image(image_name):
+    """'Abyssinian_1' -> 'Abyssinian', 'English_Cocker_Spaniel_12' -> 'English_Cocker_Spaniel'."""
+    return re.sub(r"_\d+$", "", image_name)
+
+
+def parse_annotation_file(ann_file):
+    """Return list of (image_name, class_id) skipping comment lines."""
+    entries = []
+    with open(ann_file) as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            parts = line.split()
+            image_name = parts[0]
+            class_id = int(parts[1])
+            entries.append((image_name, class_id))
+    return entries
+
+
+def main():
+    args = parse_args()
+    images_dir = args.images_dir
+    annotations_dir = args.annotations_dir
+    out_dir = args.out_dir
+
+    trainval_file = os.path.join(annotations_dir, "trainval.txt")
+    test_file = os.path.join(annotations_dir, "test.txt")
+
+    splits = [
+        (trainval_file, "train"),
+        (test_file, "val"),
+    ]
+
+    for ann_file, subset in splits:
+        if not os.path.exists(ann_file):
+            print(f"Warning: {ann_file} not found, skipping {subset} split.")
+            continue
+
+        entries = parse_annotation_file(ann_file)
+        for image_name, class_id in entries:
+            breed = class_name_from_image(image_name)
+            # Zero-pad class_id to three digits to match the NNN.ClassName convention
+            class_folder = f"{class_id:03d}.{breed}"
+
+            src = os.path.join(images_dir, image_name + ".jpg")
+            if not os.path.exists(src):
+                print(f"Warning: image not found: {src}")
+                continue
+
+            dst_dir = os.path.join(out_dir, subset, class_folder)
+            os.makedirs(dst_dir, exist_ok=True)
+
+            dst = os.path.join(dst_dir, image_name + ".jpg")
+            if not os.path.exists(dst):
+                shutil.copy2(src, dst)
+
+    for subset in ("train", "val"):
+        subset_path = os.path.join(out_dir, subset)
+        if os.path.exists(subset_path):
+            n = len(os.listdir(subset_path))
+            print(f"{subset}: {n} classes")
+
+    print(f"Output saved to: {out_dir}")
+
+
+if __name__ == "__main__":
+    main()