diff --git a/.github/workflows/check-lint.yml b/.github/workflows/check-lint.yml index 1ec2221..bea056d 100644 --- a/.github/workflows/check-lint.yml +++ b/.github/workflows/check-lint.yml @@ -20,6 +20,7 @@ jobs: uses: astral-sh/setup-uv@v1 with: version-file: "pyproject.toml" + version: latest - name: Sync Dependencies run: uv sync --dev diff --git a/.gitignore b/.gitignore index 2ae665f..fbaf140 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,11 @@ poetry.toml # LSP config files pyrightconfig.json + +# Datasets +testdata2/ +dataset/ +ds/ + +# weghts +weights/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 526eda1..e2ed766 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,10 +13,35 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Utility 'FewShotModel' [5](https://github.com/fundacaocerti/Fewpy/pull/5) - Utility 'Preprocessor' [5](https://github.com/fundacaocerti/Fewpy/pull/5) - Model 'FPTRANS' [4](https://github.com/fundacaocerti/Fewpy/pull/4) - - Model 'anomalyCLIP' [5](https://github.com/fundacaocerti/Fewpy/pull/5) + - Model 'AnomalyCLIP' [5](https://github.com/fundacaocerti/Fewpy/pull/5) - Model 'AirShot' [6](https://github.com/fundacaocerti/Fewpy/pull/6) - Model 'Qwen' [7](https://github.com/fundacaocerti/Fewpy/pull/7) + - Utility 'CVATAdapter' [9](https://github.com/fundacaocerti/Fewpy/pull/9) + - Documentation 'Fewpy/examples' [9](https://github.com/fundacaocerti/Fewpy/pull/9) + - Model 'TipAdapter' [20](https://github.com/fundacaocerti/Fewpy/pull/20) ### Changed - Utility 'FSLDataset' now supports shortest edge resizing [6](https://github.com/fundacaocerti/Fewpy/pull/6) - Utility 'FSLDataset' now supports bounding box normalization [7](https://github.com/fundacaocerti/Fewpy/pull/7) + - Model 'FPTRANS' had outputs adapted to the libraries new standard [8](https://github.com/fundacaocerti/Fewpy/pull/8) + - Model 'AnomalyCLIP' had outputs adapted to the libraries new standard [8](https://github.com/fundacaocerti/Fewpy/pull/8) + - Model 'AirShot' had outputs adapted to the libraries new standard [8](https://github.com/fundacaocerti/Fewpy/pull/8) + - Model 'Qwen' had outputs adapted to the libraries new standard [8](https://github.com/fundacaocerti/Fewpy/pull/8) + - Model 'Airshot' had its weight path search logic improved [9](https://github.com/fundacaocerti/Fewpy/pull/9) + - Model 'AnomalyCLIP' had its weight path search logic improved [9](https://github.com/fundacaocerti/Fewpy/pull/9) + - Utility 'FSLDataset' is now used directly with torch.utils.data.DataLoader [11](https://github.com/fundacaocerti/Fewpy/pull/11) + - Documentation 'Fewpy/examples' cvat exmple import typo fixed [11](https://github.com/fundacaocerti/Fewpy/pull/11) + - Model 'Airshot' now supports batched inputs (query) [11](https://github.com/fundacaocerti/Fewpy/pull/11) + - Model 'Qwen' now supports batched inputs (query) [11](https://github.com/fundacaocerti/Fewpy/pull/11) + - Model 'AnomalyCLIP' fixed prompt learner weight loading isses [11](https://github.com/fundacaocerti/Fewpy/pull/11) + - Utility 'FSLDataset' now implements multiple preprocessing routines to be used during inference and fine tuning [17](https://github.com/fundacaocerti/Fewpy/pull/17) + - Utility 'FewShotModel' improved to support model fine tuning [17](https://github.com/fundacaocerti/Fewpy/pull/17) + - Model 'AnomalyCLIP' predict method updated to support training mode [17](https://github.com/fundacaocerti/Fewpy/pull/17) + - Model 'Airshot' predict method updated to support training mode [17](https://github.com/fundacaocerti/Fewpy/pull/17) + - Model 'Qwen' predict method updated to support training mode [17](https://github.com/fundacaocerti/Fewpy/pull/17) + - Model 'Qwen' supports LoRA and Quantization [17](https://github.com/fundacaocerti/Fewpy/pull/17) + - Utility 'FSLDataset' fixed multiprocessing bug when preprocessing support set [20](https://github.com/fundacaocerti/Fewpy/pull/20) + - Utility 'FSLDataset' added new preprocessing method to be used with TipAdapter [20](https://github.com/fundacaocerti/Fewpy/pull/20) + - Utility 'FSLDataset' uses PreprocessingMethod(Enum) for preprocessing method selection [20](https://github.com/fundacaocerti/Fewpy/pull/20) + - Utility 'FSLDataset' fixed bounding box normalization bug [20](https://github.com/fundacaocerti/Fewpy/pull/20) + - Model 'FPTRANS' predict updated to support training mode [20](https://github.com/fundacaocerti/Fewpy/pull/20) \ No newline at end of file diff --git a/README.md b/README.md index 95b32ed..d9e53c4 100644 --- a/README.md +++ b/README.md @@ -31,25 +31,27 @@ Fewpy makes it easy to load and run highly performant models on benchmarks to ex * **labels**: (list[dict]) Optional annotations (bboxes/classes) for the query images. * **img_size**: (tuple/int) Target (H, W) for resizing. * **max_size**: (int) Optional maximum edge length constraint for resizing. +* **antialias**: (bool) If set True Resizing uses antialias. +* **interpolation**: (torchvision.transforms.InterpolationMode) Interpolation mode used in the resizing operation. * **pixel_norm**: (tuple) Mean and Std for normalization. -* **support_set_preprocessing_method**: (str) Strategy for support set transformation. Options include: - * `"standard"`: Basic transform and stacking. - * `"resize_annotations"`: Resizes images and mask-style `s_y`. - * `"normalize_annotations"`: Scales `s_y` bboxes to [0, 1]. - * `"detection_crop"`: Crops `s_x` around `s_y` bboxes with padding. - * `"norm_detection_crop"`: Padded crop with pixel normalization. - * `"resize_labels"`: Resizes both support annotations and query labels. - * `"none"`: Skips preprocessing. +* **center_crop**: (int) Size of the center crop. +* **support_set_preprocessing_method**: (PreprocessingMethod) Strategy for support set transformation. Options include: + * `PreprocessingMethod.STANDARD`: Basic transform and stacking. + * `PreprocessingMethod.RESIZE_SUPPORT_GT`: Resizes images and mask-style `s_y`. + * `PreprocessingMethod.NORMALIZE_SUPPORT_GT`: Scales `s_y` bboxes to [0, 1]. + * `PreprocessingMethod.DETECTION_CROP`: Crops `s_x` around `s_y` bboxes with padding. + * `PreprocessingMethod.NORM_DETECTION_CROP`: Padded crop with pixel normalization. + * `PreprocessingMethod.AUGMENT_SUPPORT_IMAGES`: Creates more support images through augmentation of the support set. + * `PreprocessingMethod.NONE`: Skips preprocessing. * **transform_datapoints**: (bool) If True, applies the transform pipeline to Query images and scales their associated `labels`. #### Key Logic 1. **Transform Pipeline**: A `torchvision.transforms.Compose` pipeline is built using `Resize`, `ToTensor`, and `Normalize` based on initialization arguments. 2. **Lazy Support Set Preprocessing**: - - On the first `__getitem__` call, the dataset executes a specific method determined by `support_set_preprocessing_method`. + - On initialization `__init__` call, the dataset executes a specific method determined by `support_set_preprocessing_method`. - **Stacking**: Images are stacked into a 4D tensor if shapes are uniform; otherwise, they remain a list of tensors. - **Detection Cropping**: Uses `_padded_crop` to zoom into support objects with a 0.5 padding ratio before resizing. -3. **Query Label Scaling**: If `transform_datapoints` is enabled and `labels` are provided, bounding boxes for query images are automatically rescaled to match the new `img_size`. -4. **Caching**: Once processed, `self.support_set_preproc` is set to True to bypass reprocessing for subsequent indices. +3. **Query Label Scaling**: If `transform_datapoints` is enabled and `labels` are provided, bounding boxes and segmentation maps for query images are automatically rescaled to match the new `img_size`. #### Utility: fsl_collate Standard collators stack every element, which can lead to redundant memory usage in FSL because the support set is usually identical for every query in a batch. `fsl_collate` optimizes this by: @@ -59,7 +61,7 @@ Standard collators stack every element, which can lead to redundant memory usage ### Usage Example ```python -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, PreprocessingMethod # Initialize with detection-based cropping for the support set ds = FSLDataset( @@ -67,7 +69,7 @@ ds = FSLDataset( s_x=s_imgs, s_y=s_annots, img_size=(320, 320), - support_set_preprocessing_method="detection_crop" + support_set_preprocessing_method=PreprocessingMethod.DETECTION_CROP ) # Use custom collate to prevent support set duplication in memory @@ -103,6 +105,7 @@ model = FewShotModel(model="AirShot", config=args) #### 1. AnomalyCLIP * **Task:** Segmentation * **Description:** Implementation of [AnomalyCLIP](https://arxiv.org/abs/2310.18961). Optimized for zero-shot or few-shot anomaly detection and segmentation. + * **Detailed Documentation:** [examples/models/anomalyclip.md](./examples/models/anomalyclip.md) #### 2. FPTRANS @@ -120,8 +123,15 @@ model = FewShotModel(model="AirShot", config=args) #### 4. Qwen * **Task:** Detection * **Description:** A wrapper for [Qwen3-VL-8B-Instruct](https://arxiv.org/abs/2505.09388) that enables few-shot prompting for object detection via Vision-Language Models. + * **Detailed Documentation:** [examples/models/qwen.md](./examples/models/qwen.md) +#### 5. TipAdapter +* **Task:** Classification +* **Description:** Based on [Tip-Adapter](https://arxiv.org/abs/2207.09519). A training-free Adaption of CLIP for Few-shot Classification. + +* **Detailed Documentation:** [examples/models/qwen.md](./examples/models/tipadapter.md) + ## CVAT Fewpy can be used with CVAT for data annotation. For that purpose inference outputs a dictionary and fewpy implements CVATAdapter. For examples of integration with cvat-cli and CVAT Community refer to [CVAT Examples](./examples/cvat/) diff --git a/examples/fine_tuning/airshot.py b/examples/fine_tuning/airshot.py index 390878f..b7a56ef 100644 --- a/examples/fine_tuning/airshot.py +++ b/examples/fine_tuning/airshot.py @@ -1,5 +1,5 @@ from fewpy.inference import FewShotModel -from fewpy.util.data import FSLDataset +from fewpy.util.data import FSLDataset, PreprocessingMethod from torch.utils.data import DataLoader from torch.optim import AdamW @@ -127,7 +127,7 @@ def fsl_collate(batch): img_size=600, max_size=1000, pixel_norm=((103.530, 116.280, 123.675), (1.0, 1.0, 1.0)), # (mean, std) - support_set_preprocessing_method="detection_crop", + support_set_preprocessing_method=PreprocessingMethod.DETECTION_CROP, ) dl = DataLoader( diff --git a/examples/fine_tuning/anomalyclip.py b/examples/fine_tuning/anomalyclip.py index eaba32a..bf422cc 100644 --- a/examples/fine_tuning/anomalyclip.py +++ b/examples/fine_tuning/anomalyclip.py @@ -1,5 +1,5 @@ from fewpy.inference import Preprocessor, FewShotModel -from fewpy.util.data import FSLDataset +from fewpy.util.data import FSLDataset, PreprocessingMethod from fewpy.util.loss import BinaryDiceLoss, FocalLoss from torch.utils.data import DataLoader @@ -71,7 +71,7 @@ def collate_fn(batch): labels=labels, img_size=(IMG_SIZE, IMG_SIZE), pixel_norm=((0.485, 0.456, 0.0406), (0.229, 0.224, 0.225)), # (mean, std) - support_set_preprocessing_method="resize_labels", + support_set_preprocessing_method=PreprocessingMethod.RESIZE_SUPPORT_GT, ) dl = DataLoader( diff --git a/examples/fine_tuning/fptrans.py b/examples/fine_tuning/fptrans.py new file mode 100644 index 0000000..1606593 --- /dev/null +++ b/examples/fine_tuning/fptrans.py @@ -0,0 +1,228 @@ +from fewpy.inference import FewShotModel +from fewpy.util.data import FSLDataset, PreprocessingMethod + +from pathlib import Path +import xml.etree.ElementTree as ET + +import PIL +import numpy as np + +import torch +from torch.optim import AdamW +from tqdm import tqdm + + +def fptrans_collate(batch): + + batch, labels, s_x, s_y = zip(*batch) + batch = torch.stack(batch) + s_x = torch.stack(s_x) + s_y = torch.stack(s_y) + + return batch, torch.stack(labels), s_x, s_y + +# prepare support and query data +K = 1 +N = 64 +CLASSES = ["bottle",] +IMG_SIZE = 700 +epochs = 4 +lr = 1e-5 +loss_fn = torch.nn.CrossEntropyLoss(ignore_index=255) + +dataset = Path("./ds").expanduser() +annotations = dataset / "Annotations" +images = dataset / "JPEGImages" +segmentation = dataset / "SegmentationClass" + +annotations = list(annotations.glob("*.xml")) +# print("retrieving", len(annotations), "annotations") +# print() + +indices = np.random.permutation(len(annotations)) +query_targets = [] +query_images = [] +support_images = [] +support_ground_truth = [] +sizes = [] + +for i in indices: + annotation = annotations[i] + # print(f"looking into {annotation}") + + tree = ET.parse(annotation) + root = tree.getroot() + image = images / root.find("filename").text + s_yi = segmentation / f"{root.find("filename").text[:-4]}.png" + if not s_yi.exists() or not image.exists(): + # print(f"Segmentation {s_yi} does not exist!\n\n") + continue + first = True + cls = None + + size = root.find("size") + width = int(size.find("width").text) + height = int(size.find("height").text) + + for obj in root.findall("object"): + name = obj.find("name").text + if name not in CLASSES: + first = True + break + + first = False + cls = name + + if first == True: + # print("no object from class list found!") + continue + + if len(support_images) < K: + mask = PIL.Image.open(s_yi).convert("L") + mask_tensor = torch.from_numpy(np.array(mask)) + mask_tensor = (mask_tensor > 0).float().unsqueeze(0) + support_images.append(PIL.Image.open(image).convert("RGB")) + support_ground_truth.append(mask_tensor) + + elif len(query_images) < N: + query_images += [PIL.Image.open(image).convert("RGB")] + mask = PIL.Image.open(s_yi).convert("L") + mask = torch.from_numpy(np.array(mask)) + mask = (mask > 0).long().unsqueeze(0) + query_targets += [mask] + sizes.append(query_images[-1].size) + + if len(support_images) >= K and len(query_images) >= N: + break + +ds = FSLDataset( + x=query_images, + s_x=support_images, + s_y=support_ground_truth, + labels=query_targets, + img_size=(IMG_SIZE, IMG_SIZE), + pixel_norm=((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), # (mean, std) + support_set_preprocessing_method=PreprocessingMethod.RESIZE_SUPPORT_GT, +) + +dl = torch.utils.data.DataLoader( + ds, + batch_size=1, + shuffle=False, + collate_fn=fptrans_collate # fptrans collate function +) + +""" +kshot: int, Number of support examples per class in the few-shot setup (e.g., 1-shot learning) +dataset: str, The dataset being used for training or evaluation (e.g., Pascal VOC) +backbone: str, The neural network backbone used for feature extraction (Vision Transformer Base) +split: int, The specific dataset split used for cross-validation/testing +checkpoint: dict | str | None, Dictionary or path containing pre-trained model weights to load +Probs_return: bool, Whether the model should return soft probability maps instead of hard logits/masks +drop_dim: int, The specific tensor dimension along which dropout is applied +drop_rate: float, The dropout probability rate used to prevent overfitting +block_size: int, The patch/block resolution for the Vision Transformer (e.g., 16x16 pixel patches) +height: int, The input image height (often paired with img_size) +pretrained: str, Path or identifier for the backbone's pre-trained weights (e.g., ImageNet weights) +SAHI: bool, Flag to enable Slicing Aided Hyper Inference (used for detecting/segmenting small objects) +bg_num: int, Number of background prototypes or background tokens used to model background context +bsz: int, The batch size used in SAHI inference +img_size: int, The overall spatial dimensions to resize the input image to (e.g., 700x700) +training: bool, Boolean flag indicating whether the network is in training mode or evaluation mode +vit_depth: int, The specific depth (number of transformer layers) extracted from the ViT backbone +vit_stride: int, The stride used for ViT patch extraction (affects the sequence length and overlap) +coco2pascal: bool, Flag for cross-dataset evaluation (e.g., training on COCO, evaluating on Pascal VOC) +num_prompt: int, The number of learnable visual prompt tokens injected into the transformer +pt_std: float, Standard deviation for the normal distribution used to initialize the prompt tokens +""" + +args = { + "kshot": 1, + "dataset": "pascal", + "backbone": "VIT-B", + "checkpoint": None, + "Probs_return": False, + "drop_dim": 1, + "drop_rate": 0.3, + "block_size": 16, + "height": IMG_SIZE, + "pretrained": "", + "SAHI": False, + "bg_num": 5, + "bsz": 32, + "img_size": IMG_SIZE, + "training": True, + "vit_depth": 10, + "vit_stride": 23, + "num_prompt": 72, + "pair_lossW": 0.02, +} + +model = FewShotModel( + model="FPTRANS", + config=args +) + +model.train() + +for param in model.parameters(): + param.requires_grad = False + +for param in model.model.purifier.parameters(): + param.requires_grad = True + +parameters = [param for param in model.parameters() if param.requires_grad] +optimizer = AdamW(parameters, lr=lr) + +""" +FPTRANS.predict: +Args: + batched_inputs: a list, batched outputs of :class:`DatasetMapper` . + Each item in the list contains the inputs for one image. + For now, each item in the list is a dict that contains: + * x: Tensor, batch of images in (B, C, H, W) format. + * s_x: Tensor, batch of support images in (B, K, C, H, W) format. + * s_y: Tensor, batch of ground truth images in (B, K, H, W) format. +Returns: + list[dict]: + Each dict is corresponds to the output of a single input image. + The dict contains a string "segmentation" under the key "task" to specify the task type, + a "data" mask, Tensor of format (H, W) and a "postproc_data" mask, Tensor of format (H, W) +""" +for epoch in range(epochs): + loss_dict = { + 'loss': [], + 'prompt': [], + 'pair': [], + } + for batch, labels, s_x, s_y in tqdm(dl): + optimizer.zero_grad() + # print(f"shapes: batch = {batch.shape}, s_x = {s_x.shape}, s_y = {s_y.shape}") + output = model.predict( + x=batch, + s_x=s_x, + s_y=s_y, + y=labels, + ) + + loss_pair = output["loss_pair"] + labels = labels.view(-1, *labels.shape[-2:]) + loss = loss_fn(output["out"], labels) + loss_prompt = loss_fn(output["out_prompt"], labels) + total_loss = loss + loss_pair + loss_prompt + + total_loss.backward() + optimizer.step() + + loss_dict["loss"].append(loss.item()) + loss_dict["pair"].append(loss_pair.item()) + loss_dict["prompt"].append(loss_prompt.item()) + + print(f"Epoch: {epoch}/{epochs}") + print(f"Pairwise Loss: {sum(loss_dict["pair"]) / len(loss_dict["pair"])}") + print(f"Prompt Loss: {sum(loss_dict["prompt"]) / len(loss_dict["prompt"])}") + print(f"Output Loss: {sum(loss_dict["loss"]) / len(loss_dict["loss"])}") + print(f"Total Loss {sum([sum(v) for v in loss_dict.values()]) / len(loss_dict["loss"])}") + +print("\nFine-tuning complete!") +torch.save(model.state_dict(), "fptrans.pth") diff --git a/examples/fine_tuning/qwen.py b/examples/fine_tuning/qwen.py index f56f42c..a676043 100644 --- a/examples/fine_tuning/qwen.py +++ b/examples/fine_tuning/qwen.py @@ -5,14 +5,12 @@ from pathlib import Path from torch.optim import AdamW -import torch.nn.functional as F import torch import xml.etree.ElementTree as ET from PIL import Image import numpy as np -import json def qwen_collate(batch): diff --git a/examples/fine_tuning/tipadapter.py b/examples/fine_tuning/tipadapter.py new file mode 100644 index 0000000..0c31ec9 --- /dev/null +++ b/examples/fine_tuning/tipadapter.py @@ -0,0 +1,170 @@ +from fewpy.inference import FewShotModel +from fewpy.util.data import FSLDataset, PreprocessingMethod, fsl_collate + +from pathlib import Path +from PIL import Image + +import xml.etree.ElementTree as ET +import numpy as np +import torch + +from torch.nn.functional import cross_entropy +from tqdm import tqdm + + +# prepare support and query data +K = 10 +n = 4 +CLASSES = ("bottle", "sofa") +epochs = 5 +batch_size = 2 +IMG_SIZE = 224 +cls2int = { + "bottle": 0, + "sofa": 1, +} +learning_rate = 1e-3 +eps = 1e-4 +epochs = 5 + +dataset = Path("./testdata2").expanduser() +annotations = dataset / "selected_annot" +images = dataset / "selected_img" + +counter = [0] * 2 + +annotations = list(annotations.glob("*.xml")) + +indices = np.random.permutation(len(annotations)) +query_targets = [] +query_images = [] +support_images = [] +support_ground_truth = [] + +for i in indices: + annotation = annotations[i] + + tree = ET.parse(annotation) + root = tree.getroot() + image = images / root.find("filename").text + cls = None + + skip = False + for obj in root.findall("object"): + name = obj.find("name").text + if name not in CLASSES or (cls is not None and name != cls): + skip = True + break + + first = False + cls = cls2int[name] + + if counter[cls] < K: + support_ground_truth.append(torch.tensor(cls).squeeze()) + support_images.append(Image.open(image).convert("RGB")) + + counter[cls] += 1 + elif len(query_images) < n: + query_images += [Image.open(image).convert("RGB")] + query_targets += [torch.tensor(cls).squeeze()] + + if all(ki >= K for ki in counter) and \ + len(query_images) >= n: + break + +""" +kshot: int, Number of examples per class in the support set +clip_model: str, Name of the CLIP model used +cache_kv: bool, Flag for KV caching +cache_dir: str, Directory to save the cached KV +augment_epoch: int, Number of epochs in the support set images augmentation process, +used to calculate the expected number of support images and transform the label tensor accordingly +load_adapter: bool, flag for loading previouly tuned adapter from a .pth file +training: bool, training mode flag +beta: float, sharpness for cache_logits calculation +alpha: float, scales clip logits when combining the logits to gen the final prediction +show_logits: bool, flag for outputing logits along with final prediction +""" +args = { + "kshot": 10, + "clip_model": "ViT-B/32", + "augment_epoch": 12, + "show_logits": True, + "training": True, + "n_novel_classes": len(CLASSES), + "cache_kv": True, + "cache_dir": "./cache" +} + +model = FewShotModel("TipAdapter", config=args) + +ds = FSLDataset( + query_images, + support_images, + support_ground_truth, + labels=query_targets, + img_size=IMG_SIZE, + center_crop=IMG_SIZE, + pixel_norm=((0.48145466, 0.48145466, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), # (mean, std)) + support_set_preprocessing_method=PreprocessingMethod.AUGMENT_SUPPORT_IMAGES, + epochs=args["augment_epoch"], +) + +dl = torch.utils.data.DataLoader( + dataset=ds, + batch_size=batch_size, + collate_fn=fsl_collate, + shuffle=False +) + +# creating optimizer +# no layer freezing, TipAdapter only returns the adapters parameters +optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, eps=eps) +T_max = epochs * len(dl) +scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max) + +""" +self.model.predict: +Args: + batched_inputs: a list, batched outputs of :class:`DatasetMapper` . + Each item in the list contains the inputs for one image. + For now, each item in the list is a dict that contains: + * x: Tensor, batch of images in (B, C, H, W) format. + * s_x: Tensor, batch of support images in (B, C, H, W) format. + * s_y: Tensor, batch of ground truth classes in (B) format. + * prompts: list[str], prompt templates for the classnames to be formatted into. + * classnames: list[str], name of all novel classes. +Returns: + list[dict]: + Each dict is corresponds to the output of a single input image. + The dict contains a string "classification" under the key "task" to specify the task type and + a "data" class prediction, Tensor of format (1) +""" +results = [] +for epoch in range(epochs): + loss_list = [] + for batch, labels, s_x, s_y in tqdm(dl): + + labels = torch.stack(labels).to(model.device) + + logits = model.predict( + x=batch, + s_x=s_x, + s_y=s_y, + classnames=CLASSES, + prompts=["an image of one or more objects of type {}"] + ) + + loss = cross_entropy(logits, labels) + loss_list.append(loss.item()) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + scheduler.step() + + current_lr = scheduler.get_last_lr()[0] + print(f"Epoch {epoch+1}/{epochs} | LR: {current_lr:.6f} | Loss: {sum(loss_list) / len(loss_list):.8f}") + +adapter_state_dict = {k: v for k, v in model.state_dict().items() if k.startswith("adapter")} +torch.save(adapter_state_dict, "tip_adapter.pth") diff --git a/examples/models/airshot.py b/examples/models/airshot.py index 99930e9..6a7cb6c 100644 --- a/examples/models/airshot.py +++ b/examples/models/airshot.py @@ -2,7 +2,7 @@ import xml.etree.ElementTree as ET from fewpy.inference import FewShotModel -from fewpy.util.data import FSLDataset, fsl_collate +from fewpy.util.data import FSLDataset, PreprocessingMethod, fsl_collate from torch.utils.data import DataLoader import torch @@ -91,7 +91,7 @@ img_size=600, max_size=1000, pixel_norm=((103.530, 116.280, 123.675), (1.0, 1.0, 1.0)), # (mean, std) - support_set_preprocessing_method="detection_crop", + support_set_preprocessing_method=PreprocessingMethod.DETECTION_CROP, ) dl = DataLoader( diff --git a/examples/models/anomalyclip.py b/examples/models/anomalyclip.py index 502db2c..9a05084 100644 --- a/examples/models/anomalyclip.py +++ b/examples/models/anomalyclip.py @@ -1,5 +1,5 @@ from fewpy.inference import Preprocessor, FewShotModel -from fewpy.util.data import FSLDataset, fsl_collate +from fewpy.util.data import FSLDataset, PreprocessingMethod, fsl_collate from torch.utils.data import DataLoader import torch @@ -48,7 +48,7 @@ s_y=support_ground_truth, img_size=(IMG_SIZE, IMG_SIZE), pixel_norm=((0.485, 0.456, 0.0406), (0.229, 0.224, 0.225)), # (mean, std) - support_set_preprocessing_method="resize_annotations", + support_set_preprocessing_method=PreprocessingMethod.RESIZE_SUPPORT_GT, ) dl = DataLoader( diff --git a/examples/models/fptrans.md b/examples/models/fptrans.md index fe4710c..f6b1649 100644 --- a/examples/models/fptrans.md +++ b/examples/models/fptrans.md @@ -55,6 +55,7 @@ args = { "coco2pascal": False, # flag for cross-dataset evaluation (e.g., training on COCO, evaluating on Pascal VOC) "num_prompt": 72, # the number of learnable visual prompt tokens injected into the transformer "pt_std": 0.02, # standard deviation for the normal distribution used to initialize the prompt tokens + "pair_lossW": 0.02 # scaling factor for pairwise loss } ``` diff --git a/examples/models/fptrans.py b/examples/models/fptrans.py index d268835..f5c5e27 100644 --- a/examples/models/fptrans.py +++ b/examples/models/fptrans.py @@ -1,5 +1,5 @@ from fewpy.inference import FewShotModel -from fewpy.util.data import FSLDataset +from fewpy.util.data import FSLDataset, PreprocessingMethod from pathlib import Path import xml.etree.ElementTree as ET @@ -95,7 +95,7 @@ def fptrans_collate(batch): s_y=support_ground_truth, img_size=(IMG_SIZE, IMG_SIZE), pixel_norm=((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), # (mean, std) - support_set_preprocessing_method="resize_annotations", + support_set_preprocessing_method=PreprocessingMethod.RESIZE_SUPPORT_GT, ) dl = torch.utils.data.DataLoader( diff --git a/examples/models/tipadapter.md b/examples/models/tipadapter.md new file mode 100644 index 0000000..d09d88e --- /dev/null +++ b/examples/models/tipadapter.md @@ -0,0 +1,73 @@ +# Running TipAdapter with fewpy + +In order to run the AirShot implementation within fewpy you will need: + - to import fewpy's FewShotModel class + - CLIP backbone weights (e.g., ViT-B/32) (optional) + - Adapter Weights in a '.pth' file (optional) + - to import fewpy's FSLDataset and PreprocessingMethod (optional) + +## Obtaining the models weights + +TipAdapter relies on CLIP weights. These can be downloaded directly through fewpy. The implementation is designed to work with standard CLIP architectures such as ViT-B/32. + +```python +from fewpy.util.download import download + +# Example for downloading ViT-B/32 weights +print(download("ViT-B/32")) +``` + +The model will download its weights automatically if not weight file is found. + +## Importing necessary fewpy tooling + +In order to instantiate the model, FewShotModel is necessary. For data handling, especially when using support image augmentation, you should import FSLDataset and PreprocessingMethod. + +```python +from fewpy.inference import FewShotModel +from fewpy.util.data import FSLDataset, PreprocessingMethod, fsl_collate +from torch.utils.data import DataLoader +``` + +## Instatiating AirShot + +To instantiate the model, you need a configuration dictionary. This dictionary controls the CLIP backbone, the k-shot settings, and whether to show logit outputs. + +An example configuration dictionary: + +```python +args = { + "kshot": 10, # Number of examples per class in the support set + "clip_model": "ViT-B/32", # Name of the CLIP model used + "cache_kv": True, # Flag for KV caching + "augment_epoch": 12, # Number of epochs for support set augmentation + "alpha": 1.0, # Scales clip logits when combining for final prediction + "beta": 1.0, # Sharpness for cache_logits calculation + "show_logits": True # Flag for outputting logits along with final prediction +} + +model = FewShotModel("TipAdapter", config=args) +``` + +## Preparing your input + +TipAdapter expects a query batch (x), a support image tensor (s_x), and a support ground truth tensor (s_y). Additionally, it requires a list of class names and prompt templates. + +```python +# The model formats classnames into the {} placeholder in the prompts +results = model.predict( + x=batch, + s_x=s_x, + s_y=s_y, + classnames=["bottle", "sofa"], + prompts=["an image of one or more objects of type {}"] +) +``` + +Notice that it is recommended to use PreprocessingMethod.AUGMENT_SUPPORT_IMAGES within the FSLDataset to properly prepare the support set according to the augment_epoch setting. + +## Model output + +In fewpy, every model has a standardized output dictionary containing a "task" value and a "data" value. For TipAdapter, the task is "classification" and the data is the predicted class tensor. If show_logits is enabled, the output also contains the final logits. + +As previouly mentioned you may or may not use fewpy's dataset class and collate function. A full example of AirShot inference using fewpy is availble in [Fewpy/examples/models/airshot.py](./tipadapter.py). diff --git a/examples/models/tipadapter.py b/examples/models/tipadapter.py new file mode 100644 index 0000000..3ba2787 --- /dev/null +++ b/examples/models/tipadapter.py @@ -0,0 +1,142 @@ +from fewpy.inference import FewShotModel +from fewpy.util.data import FSLDataset, PreprocessingMethod, fsl_collate + +from pathlib import Path +from PIL import Image + +import xml.etree.ElementTree as ET +import numpy as np +import torch +from tqdm import tqdm + + +# prepare support and query data +K = 10 +n = 4 +CLASSES = ("bottle", "sofa") +epochs = 5 +batch_size = 4 +IMG_SIZE = 224 +cls2int = { + "bottle": 0, + "sofa": 1, +} + +dataset = Path("./testdata2").expanduser() +annotations = dataset / "selected_annot" +images = dataset / "selected_img" + +counter = [0] * 2 + +annotations = list(annotations.glob("*.xml")) + +indices = np.random.permutation(len(annotations)) +query_targets = [] +query_images = [] +support_images = [] +support_ground_truth = [] + +for i in indices: + annotation = annotations[i] + + tree = ET.parse(annotation) + root = tree.getroot() + image = images / root.find("filename").text + cls = None + + skip = False + for obj in root.findall("object"): + name = obj.find("name").text + if name not in CLASSES or (cls is not None and name != cls): + skip = True + break + + first = False + cls = cls2int[name] + + if counter[cls] < K: + support_ground_truth.append(torch.tensor(cls).squeeze()) + support_images.append(Image.open(image).convert("RGB")) + + counter[cls] += 1 + elif len(query_images) < n: + query_images += [Image.open(image).convert("RGB")] + query_targets += [{"cls": torch.tensor(cls).squeeze()}] + + if all(ki >= K for ki in counter) and \ + len(query_images) >= n: + break + +""" +kshot: int, Number of examples per class in the support set +clip_model: str, Name of the CLIP model used +cache_kv: bool, Flag for KV caching +cache_dir: str, Directory to save the cached KV +augment_epoch: int, Number of epochs in the support set images augmentation process, +used to calculate the expected number of support images and transform the label tensor accordingly +load_adapter: bool, flag for loading previouly tuned adapter from a .pth file +training: bool, training mode flag +beta: float, sharpness for cache_logits calculation +alpha: float, scales clip logits when combining the logits to gen the final prediction +show_logits: bool, flag for outputing logits along with final prediction +""" +args = { + "kshot": 10, + "clip_model": "ViT-B/32", + "augment_epoch": 12, + "show_logits": True +} + +model = FewShotModel("TipAdapter", config=args) + +ds = FSLDataset( + query_images, + support_images, + support_ground_truth, + img_size=IMG_SIZE, + center_crop=IMG_SIZE, + pixel_norm=((0.48145466, 0.48145466, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), # (mean, std)) + support_set_preprocessing_method=PreprocessingMethod.AUGMENT_SUPPORT_IMAGES, + epochs=args["augment_epoch"], +) + +dl = torch.utils.data.DataLoader( + dataset=ds, + batch_size=batch_size, + collate_fn=fsl_collate, + shuffle=False +) + + +""" +self.model.predict: +Args: + batched_inputs: a list, batched outputs of :class:`DatasetMapper` . + Each item in the list contains the inputs for one image. + For now, each item in the list is a dict that contains: + * x: Tensor, batch of images in (B, C, H, W) format. + * s_x: Tensor, batch of support images in (B, C, H, W) format. + * s_y: Tensor, batch of ground truth classes in (B) format. + * prompts: list[str], prompt templates for the classnames to be formatted into. + * classnames: list[str], name of all novel classes. +Returns: + list[dict]: + Each dict is corresponds to the output of a single input image. + The dict contains a string "classification" under the key "task" to specify the task type and + a "data" class prediction, Tensor of format (1) +""" +results = [] +for batch, s_x, s_y in tqdm(dl): + results += model.predict( + x=batch, + s_x=s_x, + s_y=s_y, + classnames=CLASSES, + prompts=["an image of one or more objects of type {}"] + ) + +for i, prediction in enumerate(results): + + print(f"For object {i}, of class {query_targets[i]}:") + print(f"The model predicted class {prediction["data"]}") + print(f"Final logits: {prediction["logits"]}\n") diff --git a/examples/models/tipadapter_acc_test.py b/examples/models/tipadapter_acc_test.py new file mode 100644 index 0000000..13a707a --- /dev/null +++ b/examples/models/tipadapter_acc_test.py @@ -0,0 +1,151 @@ +from fewpy.inference import FewShotModel +from fewpy.util.data import FSLDataset, PreprocessingMethod, fsl_collate + +from pathlib import Path +from PIL import Image + +import xml.etree.ElementTree as ET +import numpy as np +import torch +from tqdm import tqdm + + +# prepare support and query data +K = 1 +n = 128 +CLASSES = ("bottle", "sofa", "bird", "dog", "cat", "horse", "sheep", "cow") +batch_size = 4 +IMG_SIZE = 224 +cls2int = {k: i for i, k in enumerate(CLASSES)} + +dataset = Path("./testdata2").expanduser() +annotations = dataset / "selected_annot" +images = dataset / "selected_img" + +counter = [0] * len(CLASSES) + +annotations = list(annotations.glob("*.xml")) + +indices = np.random.permutation(len(annotations)) +query_targets = [] +query_images = [] +support_images = [] +support_ground_truth = [] + +for i in indices: + annotation = annotations[i] + + tree = ET.parse(annotation) + root = tree.getroot() + image = images / root.find("filename").text + cls = None + + skip = False + for obj in root.findall("object"): + name = obj.find("name").text + if name not in CLASSES or (cls is not None and name != cls): + skip = True + break + + first = False + cls = cls2int[name] + + if counter[cls] < K: + support_ground_truth.append(torch.tensor(cls).squeeze()) + support_images.append(Image.open(image).convert("RGB")) + + counter[cls] += 1 + elif len(query_images) < n: + query_images += [Image.open(image).convert("RGB")] + query_targets += [{"cls": torch.tensor(cls).squeeze()}] + + if all(ki >= K for ki in counter) and \ + len(query_images) >= n: + break + +""" +kshot: int, Number of examples per class in the support set +clip_model: str, Name of the CLIP model used +cache_kv: bool, Flag for KV caching +cache_dir: str, Directory to save the cached KV +augment_epoch: int, Number of epochs in the support set images augmentation process, +used to calculate the expected number of support images and transform the label tensor accordingly +load_adapter: bool, flag for loading previouly tuned adapter from a .pth file +training: bool, training mode flag +beta: float, sharpness for cache_logits calculation +alpha: float, scales clip logits when combining the logits to gen the final prediction +show_logits: bool, flag for outputing logits along with final prediction +""" +args = { + "kshot": K, + "clip_model": "ViT-B/32", + "augment_epoch": 12, + "show_logits": True +} + +model = FewShotModel("TipAdapter", config=args) + +ds = FSLDataset( + query_images, + support_images, + support_ground_truth, + img_size=IMG_SIZE, + center_crop=IMG_SIZE, + pixel_norm=((0.48145466, 0.48145466, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), # (mean, std)) + support_set_preprocessing_method=PreprocessingMethod.AUGMENT_SUPPORT_IMAGES, + epochs=args["augment_epoch"], +) + +dl = torch.utils.data.DataLoader( + dataset=ds, + batch_size=batch_size, + collate_fn=fsl_collate, + shuffle=False +) + + +""" +self.model.predict: +Args: + batched_inputs: a list, batched outputs of :class:`DatasetMapper` . + Each item in the list contains the inputs for one image. + For now, each item in the list is a dict that contains: + * x: Tensor, batch of images in (B, C, H, W) format. + * s_x: Tensor, batch of support images in (B, C, H, W) format. + * s_y: Tensor, batch of ground truth classes in (B) format. + * prompts: list[str], prompt templates for the classnames to be formatted into. + * classnames: list[str], name of all novel classes. +Returns: + list[dict]: + Each dict is corresponds to the output of a single input image. + The dict contains a string "classification" under the key "task" to specify the task type and + a "data" class prediction, Tensor of format (1) +""" +results = [] +for batch, s_x, s_y in tqdm(dl): + results += model.predict( + x=batch, + s_x=s_x, + s_y=s_y, + classnames=CLASSES, + prompts=["an image of one or more objects of type {}"] + ) + + +class_counter = [0] * len(CLASSES) +class_correct_prediction_counter = [0] * len(CLASSES) +for i, prediction in enumerate(results): + + # print(f"For object {i}, of class {query_targets[i]}:") + # print(f"The model predicted class {prediction["data"]}") + # print(f"Final logits: {prediction["logits"]}\n") + class_counter[query_targets[i]["cls"]]+= 1 + if prediction["data"] == query_targets[i]["cls"]: + class_correct_prediction_counter[query_targets[i]["cls"]] += 1 + +for name, cls_id in cls2int.items(): + correct = class_correct_prediction_counter[cls_id] + total = class_counter[cls_id] + print(f"On class {name} of class id {cls_id} accuracy of {correct} / {total}") + +print(f"Overall acc: {sum(class_correct_prediction_counter)} / {sum(class_counter)}") diff --git a/fewpy/inference.py b/fewpy/inference.py index e6c2927..248095c 100644 --- a/fewpy/inference.py +++ b/fewpy/inference.py @@ -54,6 +54,15 @@ def __init__(self, model: str, config: dict=None, preprocessors: list[Preprocess def __load_model(model: str, config: dict): return REGISTRY[model][CONSTRUCTOR](REGISTRY[model][CONFIG](**config)).instantiate_model() + + @staticmethod + def get_available_models(): + return REGISTRY.keys() + + @property + def training(self): + + return self.model.training def get_model_in_features(self): diff --git a/fewpy/models/Qwen/qwen.py b/fewpy/models/Qwen/qwen.py index f4bef97..13c95ee 100644 --- a/fewpy/models/Qwen/qwen.py +++ b/fewpy/models/Qwen/qwen.py @@ -37,13 +37,13 @@ def eval(self): return self.train(False) - def parameters(self): + def parameters(self, recurse=True): - return self.model.parameters() + yield from self.model.parameters(recurse) - def named_parameters(self): + def named_parameters(self, prefix="", recurse=True, remove_duplicate=True): - return self.model.named_parameters() + yield from self.model.named_parameters(prefix, recurse, remove_duplicate) def predict(self, x: List[PIL.Image], s_x: List[PIL.Image]=None, s_y: List[Dict]=None, single_cls: str=None): diff --git a/fewpy/models/anomalyCLIP/anomaly_clip.py b/fewpy/models/anomalyCLIP/anomaly_clip.py index 33bdf81..762bfed 100644 --- a/fewpy/models/anomalyCLIP/anomaly_clip.py +++ b/fewpy/models/anomalyCLIP/anomaly_clip.py @@ -62,7 +62,7 @@ def __init__( design_details=None ): - super().__init__() + super(AnomalyCLIP, self).__init__() self.config = config if config is not None else AnomalyCLIPConfig() self.context_length = context_length @@ -223,7 +223,7 @@ def forward(self, image, text): logits_per_text = logits_per_image.t() - # return logits_per_image, logits_per_text + return logits_per_image, logits_per_text def train(self, mode: bool=True): super().train(mode) @@ -235,9 +235,13 @@ def train(self, mode: bool=True): return self - def parameters(self): + def parameters(self, recurse=True): - return self.prompt_learner.parameters() + yield from self.prompt_learner.parameters(recurse) + + def named_parameters(self, prefix = "", recurse = True, remove_duplicate = True): + + yield from self.prompt_learner.named_parameters(prefix, recurse, remove_duplicate) def predict(self, x: torch.Tensor, @@ -447,7 +451,6 @@ def instantiate_model(self): if not checkpoint_path.exists(): checkpoint_path = None - model = torch.jit.load(model_path, map_location=device) state_dict = model.state_dict() diff --git a/fewpy/models/anomalyCLIP/base/simple_tokenizer.py b/fewpy/models/anomalyCLIP/base/simple_tokenizer.py index 4d1b9f2..e8fe98f 100644 --- a/fewpy/models/anomalyCLIP/base/simple_tokenizer.py +++ b/fewpy/models/anomalyCLIP/base/simple_tokenizer.py @@ -1,6 +1,6 @@ import gzip import html -import os +from pathlib import Path from functools import lru_cache import ftfy @@ -9,11 +9,7 @@ @lru_cache() def default_bpe(): - return os.path.join( - os.path.dirname(os.path.abspath(__file__)), - "bpe_simple_vocab_16e6.txt.gz", - ) - + return Path(__file__).parent.parent.parent / "vocab" / "bpe_simple_vocab_16e6.txt.gz" @lru_cache() def bytes_to_unicode(): diff --git a/fewpy/models/fptrans/config.py b/fewpy/models/fptrans/config.py index 4982ebb..ab7a1b3 100644 --- a/fewpy/models/fptrans/config.py +++ b/fewpy/models/fptrans/config.py @@ -27,3 +27,4 @@ class FPTRANSConfig(BaseModel): coco2pascal: bool = Field(default=False) num_prompt: int = Field(default=72) pt_std: float = Field(default=0.02) + pair_lossW: float = Field(default=0.02) diff --git a/fewpy/models/fptrans/fptrans.py b/fewpy/models/fptrans/fptrans.py index c30f179..3631b76 100644 --- a/fewpy/models/fptrans/fptrans.py +++ b/fewpy/models/fptrans/fptrans.py @@ -18,6 +18,44 @@ interpb = partial(F.interpolate, mode='bilinear', align_corners=True) interpn = partial(F.interpolate, mode='nearest') +class PairwiseLoss(nn.Module): + def __init__(self, reduction='mean'): + super().__init__() + self.scale = 10 + self.reduction = reduction + self.loss_obj = nn.BCEWithLogitsLoss(reduction='none') + + def forward(self, x1, y1, x2, y2): + """ + x1: torch.Tensor. [B, S, C, N] + x2: torch.Tensor. [B, 1, C, N] + y1: torch.Tensor. [B, S, N], containing {0, 1, 255} + y2: torch.Tensor. [B, 1, N], containing {0, 1, 255} + """ + B, S, C, N = x1.shape + + x1 = F.normalize(x1, p=2, dim=2) # [B, S, C, N] + x1 = x1.permute(0, 1, 3, 2).reshape(B, S * N, C) # [B, N1, C] + y1 = y1.view(B, S * N, 1) # [B, N1, 1] + x2 = F.normalize(x2, p=2, dim=2) # [BQ, C, N] + x2 = x2.permute(0, 2, 1, 3).reshape(B, C, N) # [B, C, N2] + y2 = y2.view(B, 1, N) # [B, 1, N2] + + sim = torch.bmm(x1, x2) # [B, N1, N2] + lab = (y1 == y2).float() # [B, N1, N2] + ignore = y1 + y2 >= 255 + keep = ~(ignore | (y1 + y2 == 0)) + + loss_no_reduce = self.loss_obj(sim * self.scale, lab) + if self.reduction == 'none': + return loss_no_reduce, keep + elif self.reduction == 'mean': + # Compute on used positions + loss = (loss_no_reduce * keep).sum() / (keep.sum() + 1e-6) + return loss + else: + raise ValueError + class Residual(nn.Module): def __init__(self, layers, up=2): super().__init__() @@ -73,6 +111,14 @@ def __init__(self, args, logger): # Background sampler self.bg_sampler = np.random.RandomState(1289) + self.pairwise_loss = None + + def train(self, mode=True): + + if self.pairwise_loss is None: + self.pairwise_loss = PairwiseLoss() + + return super().train(mode) def build_upsampler(self, embed_dim): return Residual(nn.Sequential( @@ -114,7 +160,6 @@ def forward(self, x, s_x, s_y, y=None, out_shape=None): pair-wise loss """ - B, S, C, H, W = s_x.size() x0 = x.clone() x1 = torch.cat((s_x, x0.view(B, 1, C, H, W)), dim=1) @@ -163,18 +208,17 @@ def forward(self, x, s_x, s_y, y=None, out_shape=None): if not out_shape: out_shape = y.shape[-2:] if y is not None else (H, W) out = interpb(pred, out_shape) # [BQ, 2, *, *] - #output = dict(out=out) - - """ - if self.args.training and y is not None: + + if self.training and y is not None: # Pairwise loss + output = dict(out=out) x1 = sup_fts.flatten(3) # [B, S, C, N] y1 = sup_mask0.view(B, S, -1).long() # [B, S, N] # maybe need clone x2 = qry_fts.flatten(3) # [B, 1, C, N] y2 = interpn(y.float(), (h, w)).flatten(2).long() # [B, 1, N] - output['loss_pair'] = self.pairwise_loss(x1, y1, x2, y2) + output['loss_pair'] = self.pairwise_loss(x1, y1, x2, y2) * self.args.pair_lossW # Prompt-Proxy prediction fg_token = self.purifier(backbone_out['tokens']['fg'])[:, :, 0, 0] # [B, c] @@ -187,13 +231,18 @@ def forward(self, x, s_x, s_y, y=None, out_shape=None): output['out_prompt'] = pred_prompt return output - - else: - """ - if self.Probs_return: - return out #[bsz, 2, H, W] - else: - return out.argmax(dim=1) #[bsz, H, W] + + if not self.Probs_return: + out = out.argmax(dim=1) #[bsz, H, W] + + results = [] + for segmentation in out: + results.append({ + "task": "segmentation", + "data": segmentation, + }) + + return results def classifier(self, sup_fts, qry_fts, sup_mask): @@ -331,46 +380,6 @@ def compute_similarity(fg_proto, bg_proto, qry_fts, dist_scalar=20): pred = torch.stack((bg_distance, fg_distance), dim=1) # [B, 2, h, w] return pred - ''' - def load_weights(self, ckpt_path, device): - """ - - Parameters - ---------- - ckpt_path: Path - path to the checkpoint - strict: bool - strict mode or not - - """ - weights = torch.load(ckpt_path, map_location=device) - if "model_state" in weights: - weights = weights["model_state"] - if "state_dict" in weights: - weights = weights["state_dict"] - weights = {k.replace("module.", ""): v for k, v in weights.items()} - weights.update({k: v for k, v in self.state_dict().items() if 'original_encoder' in k}) - self.load_state_dict(weights) - ''' - - ''' - @staticmethod - def get_or_download_pretrained(backbone, progress): - if backbone not in pretrained_weights: - raise ValueError(f'Not supported backbone {backbone}. ' - f'Available backbones: {list(pretrained_weights.keys())}') - - cached_file = Path(pretrained_weights[backbone]) - if cached_file.exists(): - return cached_file - - # Try to download - url = model_urls[backbone] - cached_file.parent.mkdir(parents=True, exist_ok=True) - sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file)) - download_url_to_file(url, str(cached_file), progress=progress) - return cached_file - ''' def get_params_list(self): params = [] @@ -379,7 +388,7 @@ def get_params_list(self): params.append(var) return [{'params': params}] - def predict(self, x: torch.Tensor, s_x: torch.Tensor, s_y: torch.Tensor): + def predict(self, x: torch.Tensor, s_x: torch.Tensor, s_y: torch.Tensor, y: torch.Tensor=None): if self.args.SAHI: @@ -402,17 +411,7 @@ def predict(self, x: torch.Tensor, s_x: torch.Tensor, s_y: torch.Tensor): return novo_tensor - output = self(x, s_x, s_y) - results = [] - - for segmentation in output: - results.append({ - "task": "segmentation", - "data": segmentation, - }) - - return results - + return self(x, s_x, s_y, y) @register_constructor(name="FPTRANS", config_cls=FPTRANSConfig) class constructor_FPTRANS(): @@ -448,9 +447,11 @@ def load_weights(self, model): if "model_state" in weights: state_dict = weights["model_state"] - if "state_dict" in weights: + elif "state_dict" in weights: state_dict = weights["state_dict"] - + else: + state_dict = weights + state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} state_dict.update({k: v for k, v in model.state_dict().items() if 'original_encoder' in k}) @@ -458,6 +459,6 @@ def load_weights(self, model): if self.args.training: - return model.train() + return model.train(), device return model.eval(), device diff --git a/fewpy/models/tip-adapter/__init__.py b/fewpy/models/tip-adapter/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fewpy/models/tip-adapter/base/clip.py b/fewpy/models/tip-adapter/base/clip.py new file mode 100644 index 0000000..0e3ecdf --- /dev/null +++ b/fewpy/models/tip-adapter/base/clip.py @@ -0,0 +1,368 @@ +from collections import OrderedDict +from typing import Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1): + super().__init__() + + # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1 + self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + + self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + + self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity() + + self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + + self.relu = nn.ReLU(inplace=True) + self.downsample = None + self.stride = stride + + if stride > 1 or inplanes != planes * Bottleneck.expansion: + # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1 + self.downsample = nn.Sequential(OrderedDict([ + ("-1", nn.AvgPool2d(stride)), + ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)), + ("1", nn.BatchNorm2d(planes * self.expansion)) + ])) + + def forward(self, x: torch.Tensor): + identity = x + + out = self.relu(self.bn1(self.conv1(x))) + out = self.relu(self.bn2(self.conv2(out))) + out = self.avgpool(out) + out = self.bn3(self.conv3(out)) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + return out + + +class AttentionPool2d(nn.Module): + def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None): + super().__init__() + self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5) + self.k_proj = nn.Linear(embed_dim, embed_dim) + self.q_proj = nn.Linear(embed_dim, embed_dim) + self.v_proj = nn.Linear(embed_dim, embed_dim) + self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim) + self.num_heads = num_heads + + def forward(self, x): + x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1) # NCHW -> (HW)NC + x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC + x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC + x, _ = F.multi_head_attention_forward( + query=x, key=x, value=x, + embed_dim_to_check=x.shape[-1], + num_heads=self.num_heads, + q_proj_weight=self.q_proj.weight, + k_proj_weight=self.k_proj.weight, + v_proj_weight=self.v_proj.weight, + in_proj_weight=None, + in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]), + bias_k=None, + bias_v=None, + add_zero_attn=False, + dropout_p=0, + out_proj_weight=self.c_proj.weight, + out_proj_bias=self.c_proj.bias, + use_separate_proj_weight=True, + training=self.training, + need_weights=False + ) + + return x[0] + + +class ModifiedResNet(nn.Module): + """ + A ResNet class that is similar to torchvision's but contains the following changes: + - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool. + - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1 + - The final pooling layer is a QKV attention instead of an average pool + """ + + def __init__(self, layers, output_dim, heads, input_resolution=224, width=64): + super().__init__() + self.output_dim = output_dim + self.input_resolution = input_resolution + + # the 3-layer stem + self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(width // 2) + self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(width // 2) + self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False) + self.bn3 = nn.BatchNorm2d(width) + self.avgpool = nn.AvgPool2d(2) + self.relu = nn.ReLU(inplace=True) + + # residual layers + self._inplanes = width # this is a *mutable* variable used during construction + self.layer1 = self._make_layer(width, layers[0]) + self.layer2 = self._make_layer(width * 2, layers[1], stride=2) + self.layer3 = self._make_layer(width * 4, layers[2], stride=2) + self.layer4 = self._make_layer(width * 8, layers[3], stride=2) + + embed_dim = width * 32 # the ResNet feature dimension + self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim) + + def _make_layer(self, planes, blocks, stride=1): + layers = [Bottleneck(self._inplanes, planes, stride)] + + self._inplanes = planes * Bottleneck.expansion + for _ in range(1, blocks): + layers.append(Bottleneck(self._inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + def stem(x): + for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2), (self.conv3, self.bn3)]: + x = self.relu(bn(conv(x))) + x = self.avgpool(x) + return x + + x = x.type(self.conv1.weight.dtype) + x = stem(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + x = self.attnpool(x) + + return x + + +class LayerNorm(nn.LayerNorm): + """Subclass torch's LayerNorm to handle fp16.""" + + def forward(self, x: torch.Tensor): + orig_type = x.dtype + ret = super().forward(x.type(torch.float32)) + return ret.type(orig_type) + + +class QuickGELU(nn.Module): + def forward(self, x: torch.Tensor): + return x * torch.sigmoid(1.702 * x) + + +class ResidualAttentionBlock(nn.Module): + def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None): + super().__init__() + + self.attn = nn.MultiheadAttention(d_model, n_head) + self.ln_1 = LayerNorm(d_model) + self.mlp = nn.Sequential(OrderedDict([ + ("c_fc", nn.Linear(d_model, d_model * 4)), + ("gelu", QuickGELU()), + ("c_proj", nn.Linear(d_model * 4, d_model)) + ])) + self.ln_2 = LayerNorm(d_model) + self.attn_mask = attn_mask + + def attention(self, x: torch.Tensor): + self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None + return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0] + + def forward(self, x: torch.Tensor): + x = x + self.attention(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + +class Transformer(nn.Module): + def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None): + super().__init__() + self.width = width + self.layers = layers + self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]) + + def forward(self, x: torch.Tensor): + return self.resblocks(x) + + +class VisionTransformer(nn.Module): + def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int): + super().__init__() + self.input_resolution = input_resolution + self.output_dim = output_dim + self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False) + + scale = width ** -0.5 + self.class_embedding = nn.Parameter(scale * torch.randn(width)) + self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width)) + self.ln_pre = LayerNorm(width) + + self.transformer = Transformer(width, layers, heads) + + self.ln_post = LayerNorm(width) + self.proj = nn.Parameter(scale * torch.randn(width, output_dim)) + + def forward(self, x: torch.Tensor): + x = self.conv1(x) # shape = [*, width, grid, grid] + x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2] + x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width] + x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width] + x = x + self.positional_embedding.to(x.dtype) + x = self.ln_pre(x) + + x = x.permute(1, 0, 2) # NLD -> LND + x = self.transformer(x) + x = x.permute(1, 0, 2) # LND -> NLD + + x = self.ln_post(x[:, 0, :]) + + if self.proj is not None: + x = x @ self.proj + + return x + + +class CLIP(nn.Module): + def __init__(self, + embed_dim: int, + # vision + image_resolution: int, + vision_layers: Union[Tuple[int, int, int, int], int], + vision_width: int, + vision_patch_size: int, + # text + context_length: int, + vocab_size: int, + transformer_width: int, + transformer_heads: int, + transformer_layers: int + ): + super().__init__() + + self.context_length = context_length + + if isinstance(vision_layers, (tuple, list)): + vision_heads = vision_width * 32 // 64 + self.visual = ModifiedResNet( + layers=vision_layers, + output_dim=embed_dim, + heads=vision_heads, + input_resolution=image_resolution, + width=vision_width + ) + else: + vision_heads = vision_width // 64 + self.visual = VisionTransformer( + input_resolution=image_resolution, + patch_size=vision_patch_size, + width=vision_width, + layers=vision_layers, + heads=vision_heads, + output_dim=embed_dim + ) + + self.transformer = Transformer( + width=transformer_width, + layers=transformer_layers, + heads=transformer_heads, + attn_mask=self.build_attention_mask() + ) + + self.vocab_size = vocab_size + self.token_embedding = nn.Embedding(vocab_size, transformer_width) + self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width)) + self.ln_final = LayerNorm(transformer_width) + + self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim)) + self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) + + self.initialize_parameters() + + def initialize_parameters(self): + nn.init.normal_(self.token_embedding.weight, std=0.02) + nn.init.normal_(self.positional_embedding, std=0.01) + + if isinstance(self.visual, ModifiedResNet): + if self.visual.attnpool is not None: + std = self.visual.attnpool.c_proj.in_features ** -0.5 + nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std) + + for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]: + for name, param in resnet_block.named_parameters(): + if name.endswith("bn3.weight"): + nn.init.zeros_(param) + + proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5) + attn_std = self.transformer.width ** -0.5 + fc_std = (2 * self.transformer.width) ** -0.5 + for block in self.transformer.resblocks: + nn.init.normal_(block.attn.in_proj_weight, std=attn_std) + nn.init.normal_(block.attn.out_proj.weight, std=proj_std) + nn.init.normal_(block.mlp.c_fc.weight, std=fc_std) + nn.init.normal_(block.mlp.c_proj.weight, std=proj_std) + + if self.text_projection is not None: + nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5) + + def build_attention_mask(self): + # lazily create causal attention mask, with full attention between the vision tokens + # pytorch uses additive attention mask; fill with -inf + mask = torch.empty(self.context_length, self.context_length) + mask.fill_(float("-inf")) + mask.triu_(1) # zero out the lower diagonal + return mask + + @property + def dtype(self): + return self.visual.conv1.weight.dtype + + def encode_image(self, image): + return self.visual(image.type(self.dtype)) + + def encode_text(self, text): + x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model] + + x = x + self.positional_embedding.type(self.dtype) + x = x.permute(1, 0, 2) # NLD -> LND + x = self.transformer(x) + x = x.permute(1, 0, 2) # LND -> NLD + x = self.ln_final(x).type(self.dtype) + + # x.shape = [batch_size, n_ctx, transformer.width] + # take features from the eot embedding (eot_token is the highest number in each sequence) + x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection + + return x + + def forward(self, image, text): + image_features = self.encode_image(image) + text_features = self.encode_text(text) + + # normalized features + image_features = image_features / image_features.norm(dim=-1, keepdim=True) + text_features = text_features / text_features.norm(dim=-1, keepdim=True) + + # cosine similarity as logits + logit_scale = self.logit_scale.exp() + logits_per_image = logit_scale * image_features @ text_features.t() + logits_per_text = logits_per_image.t() + + # shape = [global_batch_size, global_batch_size] + return logits_per_image, logits_per_text diff --git a/fewpy/models/tip-adapter/base/simple_tokenizer.py b/fewpy/models/tip-adapter/base/simple_tokenizer.py new file mode 100644 index 0000000..c659cd3 --- /dev/null +++ b/fewpy/models/tip-adapter/base/simple_tokenizer.py @@ -0,0 +1,132 @@ +import gzip +import html +from pathlib import Path +from functools import lru_cache + +import ftfy +import regex as re + + +@lru_cache() +def default_bpe(): + return Path(__file__).parent.parent.parent / "vocab" / "bpe_simple_vocab_16e6.txt.gz" + + +@lru_cache() +def bytes_to_unicode(): + """ + Returns list of utf-8 byte and a corresponding list of unicode strings. + The reversible bpe codes work on unicode strings. + This means you need a large # of unicode characters in your vocab if you want to avoid UNKs. + When you're at something like a 10B token dataset you end up needing around 5K for decent coverage. + This is a signficant percentage of your normal, say, 32K bpe vocab. + To avoid that, we want lookup tables between utf-8 bytes and unicode strings. + And avoids mapping to whitespace/control characters the bpe code barfs on. + """ + bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1)) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8+n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +def get_pairs(word): + """Return set of symbol pairs in a word. + Word is represented as tuple of symbols (symbols being variable-length strings). + """ + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +class SimpleTokenizer(object): + def __init__(self, bpe_path: str = default_bpe()): + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + merges = gzip.open(bpe_path).read().decode("utf-8").split('\n') + merges = merges[1:49152-256-2+1] + merges = [tuple(merge.split()) for merge in merges] + vocab = list(bytes_to_unicode().values()) + vocab = vocab + [v+'' for v in vocab] + for merge in merges: + vocab.append(''.join(merge)) + vocab.extend(['<|startoftext|>', '<|endoftext|>']) + self.encoder = dict(zip(vocab, range(len(vocab)))) + self.decoder = {v: k for k, v in self.encoder.items()} + self.bpe_ranks = dict(zip(merges, range(len(merges)))) + self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'} + self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE) + + def bpe(self, token): + if token in self.cache: + return self.cache[token] + word = tuple(token[:-1]) + ( token[-1] + '',) + pairs = get_pairs(word) + + if not pairs: + return token+'' + + while True: + bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf'))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + new_word.extend(word[i:j]) + i = j + except: + new_word.extend(word[i:]) + break + + if word[i] == first and i < len(word)-1 and word[i+1] == second: + new_word.append(first+second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + else: + pairs = get_pairs(word) + word = ' '.join(word) + self.cache[token] = word + return word + + def encode(self, text): + bpe_tokens = [] + text = whitespace_clean(basic_clean(text)).lower() + for token in re.findall(self.pat, text): + token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8')) + bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' ')) + return bpe_tokens + + def decode(self, tokens): + text = ''.join([self.decoder[token] for token in tokens]) + text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('', ' ') + return text diff --git a/fewpy/models/tip-adapter/config.py b/fewpy/models/tip-adapter/config.py new file mode 100644 index 0000000..c196d76 --- /dev/null +++ b/fewpy/models/tip-adapter/config.py @@ -0,0 +1,16 @@ +from pydantic import BaseModel, Field +from typing import List + +class TipAdapterConfig(BaseModel): + + kshot: int = Field(default=1) + clip_model: str = Field(default="RN50") + cache_kv: bool = Field(default=False) + cache_dir: str = Field(default="") + augment_epoch: int = Field(default=1) + load_adapter: bool = Field(default=False) + training: bool = Field(default=False) + beta: float = Field(default=1.0) + alpha: float = Field(default=1.17) + show_logits: bool = Field(default=False) + n_novel_classes: int = Field(default=16) \ No newline at end of file diff --git a/fewpy/models/tip-adapter/tip-adapter.py b/fewpy/models/tip-adapter/tip-adapter.py new file mode 100644 index 0000000..60fe477 --- /dev/null +++ b/fewpy/models/tip-adapter/tip-adapter.py @@ -0,0 +1,407 @@ +from fewpy.models.register import register_constructor +from fewpy.util.download import download +from .base.clip import CLIP + +from typing import Union, List +from .config import TipAdapterConfig +from .base.simple_tokenizer import SimpleTokenizer + +import warnings + +import numpy as np +import torch.nn.functional as F +import torchvision.transforms as transforms + +import torch +import sys + +from torch import nn + +from pathlib import Path + + +_tokenizer = SimpleTokenizer() +def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> torch.LongTensor: + """ + Returns the tokenized representation of given input string(s) + + Parameters + ---------- + texts : Union[str, List[str]] + An input string or a list of input strings to tokenize + + context_length : int + The context length to use; all CLIP models use 77 as the context length + + truncate: bool + Whether to truncate the text in case its encoding is longer than the context length + + Returns + ------- + A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length] + """ + if isinstance(texts, str): + texts = [texts] + + sot_token = _tokenizer.encoder["<|startoftext|>"] + eot_token = _tokenizer.encoder["<|endoftext|>"] + all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts] + result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + + for i, tokens in enumerate(all_tokens): + if len(tokens) > context_length: + if truncate: + tokens = tokens[:context_length] + tokens[-1] = eot_token + else: + raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}") + result[i, :len(tokens)] = torch.tensor(tokens) + + return result + + +class TipAdapter(nn.Module): + + def __init__(self, config): + super(TipAdapter, self).__init__() + + self.clip = None + self.config = config + + self.adapter = None + self.first = not self.config.load_adapter + + self.alpha = self.config.alpha + self.beta = self.config.beta + + self._tranform = transforms.Compose([ + transforms.RandomResizedCrop(size=224, scale=(0.5, 1), interpolation=transforms.InterpolationMode.BICUBIC), + transforms.RandomHorizontalFlip(p=0.5), + transforms.ToTensor(), + transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)) + ]) + + def init_adapter(self, in_dim, out_dim): + + if self.adapter is None: + self.adapter = nn.Linear(in_dim, out_dim, bias=False) + + def train(self, mode=True): + + if self.adapter is not None: + self.adapter.train(mode) + + return super().train(mode) + + def parameters(self, recurse=True): + + if self.adapter is not None: + yield from self.adapter.parameters(recurse=recurse) + + def named_parameters(self, prefix='', recurse=True, remove_duplicate=True): + + if self.adapter is not None: + yield from self.adapter.named_parameters( + prefix=prefix + 'adapter.', recurse=recurse, remove_duplicate=remove_duplicate + ) + + def construct_kv(self, images, gt, num_classes): + if self.config.cache_kv: + cache_dir = Path(self.config.cache_dir) + cache_keys = cache_dir / "keys" / f"{self.config.kshot}_shots.pt" + cache_values = cache_dir / "values" / f"{self.config.kshot}_shots.pt" + + if cache_keys.exists() and cache_values.exists(): + return torch.load(cache_keys), torch.load(cache_values) + + with torch.no_grad(): + image_features = self.encode_image(images) + + num_samples = gt.shape[0] + num_aug = self.config.augment_epoch + + mismatch_msg = "Tensor size mismatch between augmented support images tensor and gt tensor" + assert image_features.shape[0] == num_aug * num_samples, mismatch_msg + + k = image_features.view(num_aug, num_samples, -1).mean(dim=0) + k /= k.norm(dim=-1, keepdim=True) + k = k.permute(1, 0) + k = k.contiguous() + v = F.one_hot(gt, num_classes=num_classes).half() + + if self.config.cache_kv: + cache_keys.parent.mkdir(parents=True, exist_ok=True) + cache_values.parent.mkdir(parents=True, exist_ok=True) + + torch.save(k, cache_keys) + torch.save(v, cache_values) + + return k, v + + def encode_text(self, prompts, classnames): + with torch.no_grad(): + clip_weights = [] + target_device = next(self.clip.parameters()).device + + for classname in classnames: + # Tokenize the prompts + classname = classname.replace('_', ' ') + texts = [t.format(classname) for t in prompts] + texts = tokenize(texts).to(target_device) + # prompt ensemble for ImageNet + class_embeddings = self.clip.encode_text(texts) + class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True) + class_embedding = class_embeddings.mean(dim=0) + class_embedding /= class_embedding.norm() + clip_weights.append(class_embedding) + + clip_weights = torch.stack(clip_weights, dim=1).to(target_device) + + return clip_weights + + def encode_image(self, x): + + with torch.no_grad(): + x = self.clip.encode_image(x) + x /= x.norm(dim=-1, keepdim=True) + + return x + + def combine_logits(self, x, clip_weights, affinity, values): + + clip_logits = 100. * x @ clip_weights + cache_logits = ((-1) * (self.beta - self.beta * affinity)).exp() @ values + final_logits = cache_logits + clip_logits * self.alpha + + return final_logits + + + def forward(self, inputs): + + x, s_x, s_y, prompts, classnames = zip(*inputs) + + # construct cache model + K, V = self.construct_kv(s_x, s_y, len(classnames)) + + # extracting textual features with clip + clip_weights = self.encode_text(prompts, classnames) + + # extract image features + x = self.encode_image(x) + + affinity = x @ K + final_logits = self.combine_logits(x, clip_weights, affinity, V) + + return final_logits + + def _initialize_adapter(self, k): + + target_device = next(self.clip.parameters()).device + self.adapter.to(device=target_device) + k = k.to(target_device) + self.adapter.half() + assert self.adapter.weight.shape == k.t().shape + with torch.no_grad(): + self.adapter.weight.copy_(k.t()) + self.adapter.train() + + self.first = False + + def predict( + self, + x: torch.Tensor, + s_x: torch.Tensor, + s_y: torch.Tensor, + prompts: list[str]=[], + classnames: list[str]=[], + ): + """ + self.model.predict: + Args: + batched_inputs: a list, batched outputs of :class:`DatasetMapper` . + Each item in the list contains the inputs for one image. + For now, each item in the list is a dict that contains: + * x: Tensor, batch of images in (B, C, H, W) format. + * s_x: Tensor, batch of support images in (B, C, H, W) format. + * s_y: Tensor, batch of ground truth classes in (B) format. + * prompts: list[str], prompt templates for the classnames to be formatted into. + * classnames: list[str], name of all novel classes. + Returns: + list[dict]: + Each dict is corresponds to the output of a single input image. + The dict contains a string "classification" under the key "task" to specify the task type and + a "data" class prediction, Tensor of format (1) + """ + # construct cache model + K, V = self.construct_kv(s_x, s_y, len(classnames)) + + # extracting textual features with clip + clip_weights = self.encode_text(prompts, classnames) + + # extract image features + x = self.encode_image(x) + + if self.training: + if self.first: + self._initialize_adapter(K) + affinity = self.adapter(x) + + return self.combine_logits(x, clip_weights, affinity, V) + + affinity = x @ K + final_logits = self.combine_logits(x, clip_weights, affinity, V) + + final_predictions = torch.argmax(final_logits, dim=1) + + results = [] + for prediction in final_predictions: + results.append({ + "data": prediction, + "task": "classification" + }) + + if self.config.show_logits: + results[-1]["logits"] = final_logits + + return results + +@register_constructor(name="TipAdapter", config_cls=TipAdapterConfig) +class constructor_TipAdapter(): + def __init__(self, args): + + self.args = args + + def instantiate_model(self): + + model = TipAdapter(self.args) + return self.load_weights(model) + + @staticmethod + def load_from_state_dict(state_dict): + + def _convert_weights_to_fp16(l): + if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)): + l.weight.data = l.weight.data.half() + if l.bias is not None: + l.bias.data = l.bias.data.half() + + if isinstance(l, nn.MultiheadAttention): + for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]: + tensor = getattr(l, attr) + if tensor is not None: + tensor.data = tensor.data.half() + + for name in ["text_projection", "proj"]: + if hasattr(l, name): + attr = getattr(l, name) + if attr is not None: + attr.data = attr.data.half() + + vit = "visual.proj" in state_dict + + if vit: + vision_width = state_dict["visual.conv1.weight"].shape[0] + vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")]) + vision_patch_size = state_dict["visual.conv1.weight"].shape[-1] + grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5) + image_resolution = vision_patch_size * grid_size + else: + counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]] + vision_layers = tuple(counts) + vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0] + output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5) + vision_patch_size = None + assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0] + image_resolution = output_width * 32 + + embed_dim = state_dict["text_projection"].shape[1] + context_length = state_dict["positional_embedding"].shape[0] + vocab_size = state_dict["token_embedding.weight"].shape[0] + transformer_width = state_dict["ln_final.weight"].shape[0] + transformer_heads = transformer_width // 64 + transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks"))) + + model = CLIP( + embed_dim, + image_resolution, vision_layers, vision_width, vision_patch_size, + context_length, vocab_size, transformer_width, transformer_heads, transformer_layers + ) + + for key in ["input_resolution", "context_length", "vocab_size"]: + if key in state_dict: + del state_dict[key] + + model.apply(_convert_weights_to_fp16) + model.load_state_dict(state_dict) + return model + + def load_weights(self, model): + + def load_to_cpu(): + warnings.warn( + "CUDA was not found or jit.load failed. Inference/Training will run on CPU and may be significantly slower.", + category=RuntimeWarning, + stacklevel=2 + ) + state_dict = torch.jit.load(model_path, map_location=device).state_dict() + clip = constructor_TipAdapter.load_from_state_dict(state_dict) + + return clip, "cpu" + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device) + + current_dir = Path(__file__).resolve().parent + model_path = current_dir / "weights" / f"tip_adapter.pt" + if not model_path.exists(): + main_dir = Path(sys.path[0]) + model_path = main_dir / "weights" / f"tip_adapter.pt" + if not model_path.exists(): + model_path = Path(download(self.args.clip_model, new_name="tip_adapter")) + + if torch.cuda.is_available(): + try: + device = 'cuda' + clip = torch.jit.load(model_path, map_location=device) + except RuntimeError: + clip, device = load_to_cpu() + else: + clip, device = load_to_cpu() + + model.clip = clip + model.clip.to(device) + model.clip.eval() + + if self.args.load_adapter: + checkpoint_path = model_path.parent / "tip_adapter.pth" + if not checkpoint_path.exists(): + raise FileNotFoundError(f"Adapter weights not found at {checkpoint_path}") + + checkpoint = torch.load(checkpoint_path) + k1, k2 = "adapter.weight", "weight" + if k1 in checkpoint.keys(): + adapter_weights = checkpoint[k1] + elif k2 in checkpoint.keys(): + adapter_weights = checkpoint[k2] + else: + e = f"Could not find adapter weights in state_dict! Weights must be under {k1} or {k2}" + raise KeyError(e) + + out_dim, in_dim = adapter_weights.shape + model.adapter = nn.Linear(in_dim, out_dim, bias=False) + model.adapter.weight = nn.Parameter(adapter_weights) + target_device = next(model.clip.parameters()).device + target_dtype = next(model.clip.parameters()).dtype + + model.adapter.to(device=target_device, dtype=target_dtype) + else: + model.init_adapter( + in_dim=1024, + out_dim=self.args.kshot * self.args.n_novel_classes, + ) + + if self.args.training: + + return model.train(), device + + return model.eval(), device diff --git a/fewpy/models/anomalyCLIP/base/bpe_simple_vocab_16e6.txt.gz b/fewpy/models/vocab/bpe_simple_vocab_16e6.txt.gz similarity index 100% rename from fewpy/models/anomalyCLIP/base/bpe_simple_vocab_16e6.txt.gz rename to fewpy/models/vocab/bpe_simple_vocab_16e6.txt.gz diff --git a/fewpy/tests/dataset_test.py b/fewpy/tests/dataset_test.py new file mode 100644 index 0000000..bc0a49d --- /dev/null +++ b/fewpy/tests/dataset_test.py @@ -0,0 +1,206 @@ +import pytest +import torch +import numpy as np +from PIL import Image + +from fewpy.util.data import d2_tensor_transform, PreprocessingMethod, FSLDataset, fsl_collate + +@pytest.fixture +def dummy_images(): + return [Image.new("RGB", (100, 100), color=(i * 10, 255, 255)) for i in range(3)] + +@pytest.fixture +def dummy_bbox_labels(): + return [{"bboxes": [[10.0, 10.0, 50.0, 50.0]], "cls": [0]} for _ in range(3)] + +@pytest.fixture +def dummy_mask_labels(): + return [torch.zeros((1, 100, 100)) for _ in range(3)] + +@pytest.fixture +def sample_data(): + imgs = [Image.new("RGB", (100, 200)) for _ in range(2)] # W=100, H=200 + labels = [ + {"bboxes": [[10, 20, 50, 80]], "cls": [1]}, # xmin, ymin, xmax, ymax + {"bboxes": [[0, 0, 10, 10]], "cls": [2]} + ] + return imgs, labels + +def test_d2_tensor_transform(): + np_img = np.zeros((100, 100, 3), dtype=np.uint8) + tensor = d2_tensor_transform(np_img) + + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (3, 100, 100) + assert tensor.dtype == torch.float32 + +def test_dataset_initialization_no_support(dummy_images): + dataset = FSLDataset(x=dummy_images, img_size=(224, 224)) + + assert len(dataset) == 3 + assert not dataset.support_set + + item = dataset[0] + assert isinstance(item, torch.Tensor) + assert item.shape == (3, 224, 224) + +def test_dataset_standard_preprocessing(dummy_images, dummy_bbox_labels): + dataset = FSLDataset( + x=dummy_images, + s_x=dummy_images, + s_y=dummy_bbox_labels, + labels=dummy_bbox_labels, + img_size=(224, 224), + support_set_preprocessing_method=PreprocessingMethod.STANDARD + ) + + assert dataset.support_set + + xi, yi, s_x, s_y = dataset[0] + + assert isinstance(xi, torch.Tensor) + assert xi.shape == (3, 224, 224) + + assert isinstance(dataset.s_x, torch.Tensor) + assert dataset.s_x.shape == (3, 3, 224, 224) + + assert "bboxes" in yi + assert isinstance(yi["bboxes"], list) + +def test_dataset_resize_support_gt(dummy_images, dummy_mask_labels): + dataset = FSLDataset( + x=dummy_images, + s_x=dummy_images, + s_y=dummy_mask_labels, + img_size=(50, 50), + support_set_preprocessing_method=PreprocessingMethod.RESIZE_SUPPORT_GT + ) + + assert isinstance(dataset.s_y, torch.Tensor) + assert dataset.s_y.shape == (3, 50, 50) + +def test_dataset_detection_crop(dummy_images, dummy_bbox_labels): + dataset = FSLDataset( + x=dummy_images, + s_x=dummy_images, + s_y=dummy_bbox_labels, + img_size=(64, 64), # Cropped size + support_set_preprocessing_method=PreprocessingMethod.DETECTION_CROP + ) + + assert isinstance(dataset.s_x, torch.Tensor) + assert dataset.s_x.shape == (3, 3, 64, 64) + + assert isinstance(dataset.s_y[0]["bboxes"], torch.Tensor) + assert "cls" in dataset.s_y[0] + +def test_dataset_augment_support_images(dummy_images, dummy_bbox_labels): + dummy_tensor_labels = [torch.tensor([1]) for _ in range(3)] + + epochs = 2 + dataset = FSLDataset( + x=dummy_images, + s_x=dummy_images, + s_y=dummy_tensor_labels, + support_set_preprocessing_method=PreprocessingMethod.AUGMENT_SUPPORT_IMAGES, + epochs=epochs + ) + + assert isinstance(dataset.s_x, torch.Tensor) + assert dataset.s_x.shape == (6, 3, 224, 224) + +def test_fsl_collate_with_query_labels(dummy_images, dummy_bbox_labels): + dataset = FSLDataset( + x=dummy_images, + s_x=dummy_images, + s_y=dummy_bbox_labels, + labels=dummy_bbox_labels, + img_size=(224, 224) + ) + + batch = [dataset[0], dataset[1]] + + batch_xi, batch_yi, s_x, s_y = fsl_collate(batch) + + assert isinstance(batch_xi, torch.Tensor) + assert batch_xi.shape == (2, 3, 224, 224) + assert isinstance(batch_yi, list) + assert len(batch_yi) == 2 + + assert s_x.shape == (3, 3, 224, 224) + +def test_fsl_collate_without_query_labels(dummy_images, dummy_bbox_labels): + dataset = FSLDataset( + x=dummy_images, + s_x=dummy_images, + s_y=dummy_bbox_labels, + img_size=(224, 224) + ) + + batch = [dataset[0], dataset[1]] + + batch_xi, s_x, s_y = fsl_collate(batch) + + assert isinstance(batch_xi, torch.Tensor) + assert batch_xi.shape == (2, 3, 224, 224) + assert s_x.shape == (3, 3, 224, 224) + +def test_normalize_support_gt(sample_data): + imgs, labels = sample_data + dataset = FSLDataset( + x=imgs, + s_x=imgs, + s_y=labels, + support_set_preprocessing_method=PreprocessingMethod.NORMALIZE_SUPPORT_GT + ) + + first_bbox = dataset.s_y[0]["bboxes"][0] + assert first_bbox == [0.1, 0.1, 0.5, 0.4] + +def test_norm_detection_crop(sample_data): + imgs, labels = sample_data + pixel_norm = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + + dataset = FSLDataset( + x=imgs, + s_x=imgs, + s_y=labels, + img_size=(128, 128), + pixel_norm=pixel_norm, + support_set_preprocessing_method=PreprocessingMethod.NORM_DETECTION_CROP + ) + + assert isinstance(dataset.s_x, torch.Tensor) + assert dataset.s_x.shape == (2, 3, 128, 128) + assert dataset.s_x.min() < 0 or dataset.s_x.max() > 1 + +def test_preprocessing_none(sample_data): + imgs, labels = sample_data + dataset = FSLDataset( + x=imgs, + s_x=imgs, + s_y=labels, + support_set_preprocessing_method=PreprocessingMethod.NONE + ) + + assert isinstance(dataset.s_x[0], Image.Image) + assert len(dataset.s_x) == 2 + +def test_unsupported_method_error(sample_data): + imgs, labels = sample_data + with pytest.raises(ValueError, match="Unsuported support_set_preprocessing_method"): + FSLDataset( + x=imgs, + s_x=imgs, + s_y=labels, + support_set_preprocessing_method="invalid_method_name" + ) + +def test_padded_crop_logic(sample_data): + imgs, labels = sample_data + dataset = FSLDataset(x=imgs) + + patch, bbox = dataset._padded_crop(imgs[0], labels[0]["bboxes"][0], output_size=(64, 64)) + + assert patch.shape == (3, 64, 64) + assert len(bbox) == 4 diff --git a/fewpy/util/data.py b/fewpy/util/data.py index 35e1a76..daf26d9 100644 --- a/fewpy/util/data.py +++ b/fewpy/util/data.py @@ -4,6 +4,7 @@ from torch.utils.data import Dataset, DataLoader from torchvision import transforms as T import torchvision.transforms.functional as F +from enum import Enum def d2_tensor_transform(image): @@ -12,28 +13,51 @@ def d2_tensor_transform(image): return torch.as_tensor(image.astype("float32").transpose(2, 0, 1)) + +class PreprocessingMethod(Enum): + + NONE = "none" + STANDARD = "standard" + NORMALIZE_SUPPORT_GT = "normalize_support_gt" + RESIZE_SUPPORT_GT = "resize_support_gt" + DETECTION_CROP = "detection_crop" + NORM_DETECTION_CROP = "norm_detection_crop" + AUGMENT_SUPPORT_IMAGES = "augment_support_images" + + class FSLDataset(Dataset): def __init__(self, x: list[Image], s_x: list[Image]=None, - s_y: list[dict] | list[torch.Tensor] | torch.Tensor=None, - labels: list[dict]=None, + s_y: list[dict] | list[torch.Tensor]=None, + labels: list[dict] | list[torch.Tensor]=None, img_size: tuple[int] | int=None, max_size: int=None, + antialias: bool=True, + interpolation: T.InterpolationMode=T.InterpolationMode.BICUBIC, pixel_norm: tuple=None, - support_set_preprocessing_method: str="standard", + center_crop: int=None, + support_set_preprocessing_method: PreprocessingMethod=PreprocessingMethod.STANDARD, transform_datapoints: bool=True, + **kwargs ) -> None: super().__init__() # set transform composition that turns pillow image objects into datapoints compatible with the library transfs = [] if img_size is not None: - if max_size is not None: - transfs.append(T.Resize(size=img_size, max_size=max_size)) - else: - transfs.append(T.Resize(img_size)) + transfs.append( + T.Resize( + size=img_size, + max_size=max_size, + antialias=antialias, + interpolation=interpolation + ) + ) + + if center_crop is not None and center_crop != 0: + transfs.append(T.CenterCrop(center_crop)) transfs.append(T.ToTensor()) if pixel_norm is not None: @@ -45,20 +69,44 @@ def __init__(self, self.support_set = (not s_x is None) and (not s_y is None) self.s_x = s_x self.s_y = s_y - self.support_set_preproc = False self.transform_datapoints = transform_datapoints self.img_size = img_size - self.method = support_set_preprocessing_method.lower() + self.method = support_set_preprocessing_method self.labels = labels self.pixel_norm = pixel_norm + if not self.support_set: + self.method = PreprocessingMethod.NONE + + match self.method: + + case PreprocessingMethod.STANDARD: + self.transform_s_x() + + case PreprocessingMethod.RESIZE_SUPPORT_GT: + self.resize_support_gt() + + case PreprocessingMethod.NORMALIZE_SUPPORT_GT: + self.normalize_annotations() + + case PreprocessingMethod.DETECTION_CROP: + self.detection_crop() + + case PreprocessingMethod.NORM_DETECTION_CROP: + self.detection_crop(True) + + case PreprocessingMethod.AUGMENT_SUPPORT_IMAGES: + epochs = kwargs["epochs"] if "epochs" in kwargs else 4 + self.augment_support_images(epochs) + + case PreprocessingMethod.NONE: + pass + + case _: + raise ValueError("Unsuported support_set_preprocessing_method configured!") + def __len__(self) -> int: return len(self.data) - - def _check_support_set(self): - - if self.s_x is None or self.s_y is None: - raise ValueError("Operation requires support set and s_x or s_y are None!") def _stack_sx(self, s_x): @@ -72,8 +120,6 @@ def _stack_sx(self, s_x): else: self.s_x = s_x - self.support_set_preproc = True - def _padded_crop(self, image, bbox, output_size=(320, 320), norm=False, padding_ratio=0.5): if isinstance(output_size, int): output_size = (output_size, output_size) @@ -115,44 +161,41 @@ def _padded_crop(self, image, bbox, output_size=(320, 320), norm=False, padding_ return support_patch, bbox - def _rescale_bboxes(self, img, bboxes=None): - - if bboxes is not None: - new_bboxes = [] - for bbox in bboxes: - new_bboxes.append([ - bbox[0] * ratio_w, - bbox[1] * ratio_h, - bbox[2] * ratio_w, - bbox[3] * ratio_h, - ]) - - return bboxes + def _transform(self, img, gt): + + W, H = img.size + xi = self.transf(img) + if isinstance(gt, dict) and "bbboxes" in gt.keys(): + new_gt = {k: gt[k] for k in gt.keys()} + new_gt["bboxes"] = list() + w, h = xi.shape[-1], xi.shape[-2] + bboxes = torch.tensor(gt["bboxes"]) + scale = torch.tensor([w / W, h / H, w / W, h / H]) + new_gt["bboxes"] = (bboxes * scale).tolist() + gt = new_gt + elif isinstance(gt, torch.Tensor) and len(gt.shape) > 2: + new_gt = T.functional.resize( + gt, self.img_size, + interpolation=T.functional.InterpolationMode.NEAREST, + ) + gt = new_gt + + return xi, gt def transform_s_x(self): - self._check_support_set() - s_x = [] - if "bboxes" in self._s_y.keys(): - s_y = [] - for img, gt in zip(self.s_x, self.s_y): - w, h = img.size - xi = self.transf(img) - ratio_h, ratio_w = xi.shape[-2] / h, xi.shape[-1] / w - img, bboxes = self._rescale_bboxes(ratio_h, ratio_w, gt["bbox"]) - new_gt = {k: gt[k] for k in gt.keys()} - new_gt["bboxes"] = bboxes - s_y.append(new_gt) - self.s_y = s_y - else: - for img in self.s_x: - img = self._transform(img) - s_x.append(img) + s_y = [] + + for img, gt in zip(self.s_x, self.s_y): + img, gt = self._transform(img, gt) + s_x.append(img) + s_y.append(gt) self._stack_sx(s_x) + self.s_y = s_y - def resize_annotations(self): + def resize_support_gt(self): if self.img_size is None: raise ValueError("Support Set cannot be resized if img_size is None!") @@ -169,36 +212,15 @@ def resize_annotations(self): )) self.s_y = torch.stack(new_s_y).squeeze(1) - - def resize_labels(self): - - self.resize_annotations() - - if self.labels is None: - raise ValueError("Cannot resize None!") - - if isinstance(self.labels[0], torch.Tensor): - new_labels = [] - for yi in self.labels: - new_labels.append(T.functional.resize( - yi, - self.img_size, - interpolation=T.functional.InterpolationMode.NEAREST, - )) - self.labels = new_labels - def normalize_annotations(self): - self._check_support_set() - s_x = [] s_y = [] for img, annot in zip(self.s_x, self.s_y): if isinstance(img, Image): + old_w, old_h = img.size img = self.transf(img) - old_h, old_w = img.size - if "bboxes" not in annot.keys(): raise KeyError("Bounding Box annotations should be named bboxes and be a list of bounding boxes [xmin, ymin, xmax, ymax]") @@ -238,58 +260,37 @@ def detection_crop(self, norm=False): self.s_x = torch.stack(s_x) self.s_y = s_y - self.support_set_preproc = True + def augment_support_images(self, epochs): + + _transform = T.Compose([ + T.RandomResizedCrop(size=224, scale=(0.5, 1), interpolation=T.InterpolationMode.BICUBIC), + T.RandomHorizontalFlip(p=0.5), + T.ToTensor(), + T.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)) + ]) + + augmented_set = [] + for _ in range(epochs): + for img in self.s_x: + augmented_set.append(_transform(img)) + + self.s_x = torch.stack(augmented_set) + self.s_y = torch.stack(self.s_y) def __getitem__(self, index: int): xi = self.data[index] - yi = dict() + yi = None if self.transform_datapoints: - w, h = xi.size - xi = self.transf(xi) if self.labels is not None: - yi = {k: self.labels[index] for k in self.labels[index].keys()} - ratio_h, ratio_w = xi.shape[-2] / h, xi.shape[-1] / w - bboxes = self._rescale_bboxes(ratio_h, ratio_w, bboxes=yi["bboxes"]) - yi["bboxes"] = bboxes + xi, yi = self._transform(xi, self.labels[index]) + else: + xi = self.transf(xi) if not self.support_set: return xi - if self.support_set_preproc: - - if not yi is None: - return xi, yi, self.s_x, self.s_y - - return xi, self.s_x, self.s_y - - match self.method: - - case "standard": - self.transform_s_x() - - case "resize_annotations": - self.resize_annotations() - - case "normalize_annotations": - self.normalize_annotations() - - case "detection_crop": - self.detection_crop() - - case "norm_detection_crop": - self.detection_crop(True) - - case "resize_labels": - self.resize_labels() - - case "none": - self.support_set_preproc = True - - case _: - raise ValueError("Unsuported support_set_preprocessing_method configured!") - - if yi is not None: + if not yi is None: return xi, yi, self.s_x, self.s_y return xi, self.s_x, self.s_y diff --git a/fewpy/util/download.py b/fewpy/util/download.py index 7e5a89b..a943275 100644 --- a/fewpy/util/download.py +++ b/fewpy/util/download.py @@ -23,7 +23,8 @@ def download( model: str, cache_dir: Union[str, None] = None, -): + new_name: str="", +): url = model2url[model] @@ -32,6 +33,7 @@ def download( cache_dir = current_dir / "weights" else: cache_dir = Path(cache_dir) + cache_dir.mkdir(parents=True, exist_ok=True) filename = Path(url).name @@ -88,5 +90,8 @@ def download( raise RuntimeError( "Model has been downloaded but the SHA256 checksum does not not match" ) + + if new_name != "": + download_target = download_target.rename(f"{download_target.parent / new_name}{download_target.suffix}") return str(download_target) diff --git a/pyproject.toml b/pyproject.toml index d724b91..342a7c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ dependencies = [ "dropblock>=0.3.0", "peft>=0.18.1", "bitsandbytes>=0.49.2", + "pytest>=9.0.2", ] [dependency-groups] @@ -64,6 +65,9 @@ poe = [ "poethepoet>=0.36.0", ] +[tool.uv] +required-version = ">=0.9.5" + [tool.ruff] line-length = 80 lint.extend-select = ["I", "W", "E"] diff --git a/uv.lock b/uv.lock index 41b515e..7830ac5 100644 --- a/uv.lock +++ b/uv.lock @@ -427,6 +427,7 @@ dependencies = [ { name = "peft" }, { name = "pillow" }, { name = "pydantic" }, + { name = "pytest" }, { name = "qwen-vl-utils" }, { name = "regex" }, { name = "scikit-image" }, @@ -463,6 +464,7 @@ requires-dist = [ { name = "peft", specifier = ">=0.18.1" }, { name = "pillow", specifier = ">=10.0.0" }, { name = "pydantic", specifier = ">=2.12.3" }, + { name = "pytest", specifier = ">=9.0.2" }, { name = "qwen-vl-utils", specifier = ">=0.0.14" }, { name = "regex", specifier = ">=2025.11.3" }, { name = "scikit-image", specifier = ">=0.24.0" }, @@ -721,6 +723,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/fe/301e0936b79bcab4cacc7548bf2853fc28dced0a578bab1f7ef53c9aa75b/imageio-2.37.2-py3-none-any.whl", hash = "sha256:ad9adfb20335d718c03de457358ed69f141021a333c40a53e57273d8a5bd0b9b", size = 317646, upload-time = "2025-11-04T14:29:37.948Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "iopath" version = "0.1.9" @@ -1492,6 +1503,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fe/39/979e8e21520d4e47a0bbe349e2713c0aac6f3d853d0e5b34d76206c439aa/platformdirs-4.3.8-py3-none-any.whl", hash = "sha256:ff7059bb7eb1179e2685604f4aaf157cfd9535242bd23742eadc3c13542139b4", size = 18567, upload-time = "2025-05-07T22:47:40.376Z" }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "poethepoet" version = "0.36.0" @@ -1716,6 +1736,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/f7/925f65d930802e3ea2eb4d5afa4cb8730c8dc0d2cb89a59dc4ed2fcb2d74/pydantic_core-2.41.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c173ddcd86afd2535e2b695217e82191580663a1d1928239f877f5a1649ef39f", size = 2147775, upload-time = "2025-10-14T10:23:45.406Z" }, ] +[[package]] +name = "pygments" +version = "2.19.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631, upload-time = "2025-06-21T13:39:12.283Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" }, +] + [[package]] name = "pyparsing" version = "3.2.5" @@ -1725,6 +1754,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/5e/1aa9a93198c6b64513c9d7752de7422c06402de6600a8767da1524f9570b/pyparsing-3.2.5-py3-none-any.whl", hash = "sha256:e38a4f02064cf41fe6593d328d0512495ad1f3d8a91c4f73fc401b3079a59a5e", size = 113890, upload-time = "2025-09-21T04:11:04.117Z" }, ] +[[package]] +name = "pytest" +version = "9.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0"