🐣 born

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/visualization.jpg filter=lfs diff=lfs merge=lfs -text

configs/reduction_16.json

@@ -0,0 +1,33 @@
+{
+    "8":{
+        "qnrf": {
+            "bins": {
+                "fine":[
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, "inf"]
+                ],
+                "dynamic": [
+                    [0, 0], [1, 1], [2, 2], [3, 3],
+                    [4, 5], [6, 7], [8, "inf"]
+                ],
+                "coarse": [
+                    [0, 0], [1, 2], [3, 4], [5, 6], [7, "inf"]
+                ]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4, 5, 6, 7, 8],
+                    "average": [0, 1, 2, 3, 4, 5, 6, 7, 9.23349]
+                },
+                "dynamic": {
+                    "middle": [0, 1, 2, 3, 4.5, 6.5, 8],
+                    "average": [0, 1, 2, 3, 4.29278, 6.31441, 9.23349]
+                },
+                "coarse": {
+                    "middle": [0, 1.5, 3.5, 5.5, 7],
+                    "average": [0, 1.14978, 3.27641, 5.30609, 8.11466]
+                }
+            }
+        }
+    }
+}

configs/reduction_32.json

@@ -0,0 +1,56 @@
+{
+    "19": {
+        "qnrf": {
+            "bins": {
+                "fine": [
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, 8], [9, 9],
+                    [10, 10], [11, 11], [12, 12], [13, 13], [14, 14],
+                    [15, 15], [16, 16], [17, 17], [18, 18], [19, "inf"]
+                ],
+                "dynamic": [
+                    [0, 0], [1, 1], [2, 2], [3, 3], [4, 4],
+                    [5, 5], [6, 6], [7, 7], [8, 8], [9, 9],
+                    [10, 11], [12, 13], [14, 15], [16, 17], [18, "inf"]
+                ],
+                "coarse": [
+                    [0, 0], [1, 2], [3, 4], [5, 6], [7, 8],
+                    [9, 10], [11, 12], [13, 14], [15, 16], [17, 18],
+                    [19, "inf"]
+                ]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+                        11, 12, 13, 14, 15, 16, 17, 18, 19
+                    ],
+                    "average": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+                        11, 12, 13, 14, 15, 16, 17, 18, 23.01897
+                    ]
+                },
+                "dynamic": {
+                    "middle": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10.5,
+                        12.5, 14.5, 16.5, 18
+                    ],
+                    "average": [
+                        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10.42903,
+                        12.43320, 14.43341, 16.43521, 21.93548
+                    ]
+                },
+                "coarse": {
+                    "middle": [
+                        0, 1.5, 3.5, 5.5, 7.5, 9.5,
+                        11.5, 13.5, 15.5, 17.5, 19
+                    ],
+                    "average": [
+                        0, 1.23498, 3.36108, 5.40298, 7.41406, 9.42356,
+                        11.43094, 13.43244, 15.43697, 17.43759, 23.01897
+                    ]
+                }
+            }
+        }
+    }
+}

configs/reduction_8.json

@@ -0,0 +1,129 @@
+{   
+    "2": {
+        "sha": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.24479]
+                }
+            }
+        },
+        "shb": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.15171]
+                }
+            }
+        },
+        "nwpu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.10737]
+                }
+            }
+        },
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.09296]
+                }
+            }
+        },
+        "jhu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2],
+                    "average": [0, 1, 2.18589]
+                }
+            }
+        }
+    },
+    "4": {
+        "sha": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.29992]
+                }
+            }
+        },
+        "shb": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.41009]
+                }
+            }
+        },
+        "nwpu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.21931]
+                }
+            }
+        },
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.21937]
+                }
+            }
+        },
+        "jhu": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4],
+                    "average": [0, 1, 2, 3, 4.24058]
+                }
+            }
+        }
+    },
+    "11": {
+        "qnrf": {
+            "bins": {
+                "fine": [[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7], [8, 8], [9, 9], [10, 10], [11, "inf"]]
+            },
+            "anchor_points": {
+                "fine": {
+                    "middle": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
+                    "average": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+                }
+            }
+        }
+    }
+}

datasets/__init__.py

@@ -0,0 +1,12 @@
+from .crowd import Crowd, available_datasets, standardize_dataset_name, NWPUTest
+from .transforms import RandomCrop, Resize, RandomResizedCrop, RandomHorizontalFlip, Resize2Multiple, ZeroPad2Multiple
+from .transforms import ColorJitter, RandomGrayscale, GaussianBlur, RandomApply, PepperSaltNoise
+from .utils import collate_fn
+
+
+__all__ = [
+    "Crowd", "available_datasets", "standardize_dataset_name", "NWPUTest",
+    "RandomCrop", "Resize", "RandomResizedCrop", "RandomHorizontalFlip", "Resize2Multiple", "ZeroPad2Multiple",
+    "ColorJitter", "RandomGrayscale", "GaussianBlur", "RandomApply", "PepperSaltNoise",
+    "collate_fn",
+]

datasets/crowd.py

@@ -0,0 +1,233 @@
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset
+from torchvision.transforms import ToTensor, Normalize
+import os
+from glob import glob
+from PIL import Image
+import numpy as np
+from typing import Optional, Callable, Union, Tuple
+
+from .utils import get_id, generate_density_map
+
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+
+available_datasets = [
+    "shanghaitech_a", "sha",
+    "shanghaitech_b", "shb",
+    "ucf_qnrf", "qnrf", "ucf-qnrf",
+    "nwpu", "nwpu_crowd", "nwpu-crowd",
+    "jhu", "jhu_crowd", "jhu_crowd_v2"
+]
+
+
+def standardize_dataset_name(dataset: str) -> str:
+    assert dataset.lower() in available_datasets, f"Dataset {dataset} is not available."
+    if dataset.lower() in ["shanghaitech_a", "sha"]:
+        return "sha"
+    elif dataset.lower() in ["shanghaitech_b", "shb"]:
+        return "shb"
+    elif dataset.lower() in ["ucf_qnrf", "qnrf", "ucf-qnrf"]:
+        return "qnrf"
+    elif dataset.lower() in ["nwpu", "nwpu_crowd", "nwpu-crowd"]:
+        return "nwpu"
+    else:  # dataset.lower() in ["jhu", "jhu_crowd", "jhu_crowd_v2"]
+        return "jhu"
+
+
+class Crowd(Dataset):
+    def __init__(
+        self,
+        dataset: str,
+        split: str,
+        transforms: Optional[Callable] = None,
+        sigma: Optional[float] = None,
+        return_filename: bool = False,
+        num_crops: int = 1,
+    ) -> None:
+        """
+        Dataset for crowd counting.
+        """
+        assert dataset.lower() in available_datasets, f"Dataset {dataset} is not available."
+        assert split in ["train", "val"], f"Split {split} is not available."
+        assert num_crops > 0, f"num_crops should be positive, got {num_crops}."
+
+        self.dataset = standardize_dataset_name(dataset)
+        self.split = split
+
+        self.__find_root__()
+        self.__make_dataset__()
+        self.__check_sanity__()
+        self.indices = list(range(len(self.image_names)))
+
+        self.to_tensor = ToTensor()
+        self.normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        self.transforms = transforms
+
+        self.sigma = sigma
+        self.return_filename = return_filename
+        self.num_crops = num_crops
+
+    def __find_root__(self) -> None:
+        # if self.dataset == "sha":
+        #     self.root = os.path.join(curr_dir, "..", "data", "ShanghaiTech_A")
+        # elif self.dataset == "shb":
+        #     self.root = os.path.join(curr_dir, "..", "data", "ShanghaiTech_B")
+        # elif self.dataset == "qnrf":
+        #     self.root = os.path.join(curr_dir, "..", "data", "QNRF")
+        # elif self.dataset == "nwpu":
+        #     self.root = os.path.join(curr_dir, "..", "data", "NWPU")
+        # else:  # self.dataset == "jhu"
+        #     self.root = os.path.join(curr_dir, "..", "data", "JHU")
+        self.root = os.path.join(curr_dir, "..", "data", self.dataset)
+
+    def __make_dataset__(self) -> None:
+        image_npys = glob(os.path.join(self.root, self.split, "images", "*.npy"))
+        if len(image_npys) > 0:
+            self.image_type = "npy"
+            image_names = image_npys
+        else:
+            self.image_type = "jpg"
+            image_names = glob(os.path.join(self.root, self.split, "images", "*.jpg"))
+
+        label_names = glob(os.path.join(self.root, self.split, "labels", "*.npy"))
+        image_names = [os.path.basename(image_name) for image_name in image_names]
+        label_names = [os.path.basename(label_name) for label_name in label_names]
+        image_names.sort(key=get_id)
+        label_names.sort(key=get_id)
+        image_ids = tuple([get_id(image_name) for image_name in image_names])
+        label_ids = tuple([get_id(label_name) for label_name in label_names])
+        assert image_ids == label_ids, "image_ids and label_ids do not match."
+        self.image_names = tuple(image_names)
+        self.label_names = tuple(label_names)
+
+    def __check_sanity__(self) -> None:
+        if self.dataset == "sha":
+            if self.split == "train":
+                assert len(self.image_names) == len(self.label_names) == 300, f"ShanghaiTech_A train split should have 300 images, but found {len(self.image_names)}."
+            else:
+                assert len(self.image_names) == len(self.label_names) == 182, f"ShanghaiTech_A val split should have 182 images, but found {len(self.image_names)}."
+        elif self.dataset == "shb":
+            if self.split == "train":
+                assert len(self.image_names) == len(self.label_names) == 400, f"ShanghaiTech_B train split should have 400 images, but found {len(self.image_names)}."
+            else:
+                assert len(self.image_names) == len(self.label_names) == 316, f"ShanghaiTech_B val split should have 316 images, but found {len(self.image_names)}."
+        elif self.dataset == "nwpu":
+            if self.split == "train":
+                assert len(self.image_names) == len(self.label_names) == 3109, f"NWPU train split should have 3109 images, but found {len(self.image_names)}."
+            else:
+                assert len(self.image_names) == len(self.label_names) == 500, f"NWPU val split should have 500 images, but found {len(self.image_names)}."
+        elif self.dataset == "qnrf":
+            if self.split == "train":
+                assert len(self.image_names) == len(self.label_names) == 1201, f"UCF_QNRF train split should have 1201 images, but found {len(self.image_names)}."
+            else:
+                assert len(self.image_names) == len(self.label_names) == 334, f"UCF_QNRF val split should have 334 images, but found {len(self.image_names)}."
+        else:  # self.dataset == "jhu"
+            if self.split == "train":
+                assert len(self.image_names) == len(self.label_names) == 2772, f"JHU train split should have 2772 images, but found {len(self.image_names)}."
+            else:
+                assert len(self.image_names) == len(self.label_names) == 1600, f"JHU val split should have 1600 images, but found {len(self.image_names)}."
+
+    def __len__(self) -> int:
+        return len(self.image_names)
+    
+    def __getitem__(self, idx: int) -> Union[Tuple[Tensor, Tensor, Tensor], Tuple[Tensor, Tensor, Tensor, str]]:
+        image_name = self.image_names[idx]
+        label_name = self.label_names[idx]
+
+        image_path = os.path.join(self.root, self.split, "images", image_name)
+        label_path = os.path.join(self.root, self.split, "labels", label_name)
+
+        if self.image_type == "npy":
+            with open(image_path, "rb") as f:
+                image = np.load(f)
+            image = torch.from_numpy(image).float() / 255.  # normalize to [0, 1]
+        else:
+            with open(image_path, "rb") as f:
+                image = Image.open(f).convert("RGB")
+            image = self.to_tensor(image)
+
+        with open(label_path, "rb") as f:
+            label = np.load(f)
+
+        label = torch.from_numpy(label).float()
+
+        if self.transforms is not None:
+            images_labels = [self.transforms(image.clone(), label.clone()) for _ in range(self.num_crops)]
+            images, labels = zip(*images_labels)
+        else:
+            images = [image.clone() for _ in range(self.num_crops)]
+            labels = [label.clone() for _ in range(self.num_crops)]
+
+        images = [self.normalize(img) for img in images]
+        if idx in self.indices:
+            density_maps = torch.stack([generate_density_map(label, image.shape[-2], image.shape[-1], sigma=self.sigma) for image, label in zip(images, labels)], 0)
+        else:
+            labels = None
+            density_maps = None
+
+        image_names = [image_name] * len(images)
+        images = torch.stack(images, 0)
+
+        if self.return_filename:
+            return images, labels, density_maps, image_names
+        else:
+            return images, labels, density_maps
+
+
+class NWPUTest(Dataset):
+    def __init__(
+        self,
+        transforms: Optional[Callable] = None,
+        sigma: Optional[float] = None,
+        return_filename: bool = False,
+    ) -> None:
+        """
+        The test set of NWPU-Crowd dataset. The test set is not labeled, so only images are returned.
+        """
+        self.root = os.path.join(curr_dir, "..", "data", "nwpu")
+
+        image_npys = glob(os.path.join(self.root, "test", "images", "*.npy"))
+        if len(image_npys) > 0:
+            self.image_type = "npy"
+            image_names = image_npys
+        else:
+            self.image_type = "jpg"
+            image_names = glob(os.path.join(self.root, "test", "images", "*.jpg"))
+
+        image_names = [os.path.basename(image_name) for image_name in image_names]
+        assert len(image_names) == 1500, f"NWPU test split should have 1500 images, but found {len(image_names)}."
+        image_names.sort(key=get_id)
+        self.image_names = tuple(image_names)
+
+        self.to_tensor = ToTensor()
+        self.normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        self.transforms = transforms
+
+        self.sigma = sigma
+        self.return_filename = return_filename
+
+    def __len__(self) -> int:
+        return len(self.image_names)
+    
+    def __getitem__(self, idx: int) -> Union[Tensor, Tuple[Tensor, str]]:
+        image_name = self.image_names[idx]
+        image_path = os.path.join(self.root, "test", "images", image_name)
+
+        if self.image_type == "npy":
+            with open(image_path, "rb") as f:
+                image = np.load(f)
+            image = torch.from_numpy(image).float() / 255.
+        else:
+            with open(image_path, "rb") as f:
+                image = Image.open(f).convert("RGB")
+            image = self.to_tensor(image)
+        
+        label = torch.tensor([], dtype=torch.float)  # dummy label
+        image, _ = self.transforms(image, label) if self.transforms is not None else (image, label)
+        image = self.normalize(image)
+
+        if self.return_filename:
+            return image, image_name
+        else:
+            return image

datasets/transforms.py

@@ -0,0 +1,262 @@
+import torch
+from torch import Tensor
+from torchvision.transforms import ColorJitter as _ColorJitter
+import torchvision.transforms.functional as TF
+import numpy as np
+from typing import Tuple, Union, Optional, Callable
+
+
+def _crop(
+    image: Tensor,
+    label: Tensor,
+    top: int,
+    left: int,
+    height: int,
+    width: int,
+) -> Tuple[Tensor, Tensor]:
+    image = TF.crop(image, top, left, height, width)
+    if len(label) > 0:
+        label[:, 0] -= left
+        label[:, 1] -= top
+        label_mask = (label[:, 0] >= 0) & (label[:, 0] < width) & (label[:, 1] >= 0) & (label[:, 1] < height)
+        label = label[label_mask]
+
+    return image, label
+
+
+def _resize(
+    image: Tensor,
+    label: Tensor,
+    height: int,
+    width: int,
+) -> Tuple[Tensor, Tensor]:
+    image_height, image_width = image.shape[-2:]
+    image = TF.resize(image, (height, width), interpolation=TF.InterpolationMode.BICUBIC, antialias=True) if (image_height != height or image_width != width) else image
+    if len(label) > 0 and (image_height != height or image_width != width):
+        label[:, 0] = label[:, 0] * width / image_width
+        label[:, 1] = label[:, 1] * height / image_height
+        label[:, 0] = label[:, 0].clamp(min=0, max=width - 1)
+        label[:, 1] = label[:, 1].clamp(min=0, max=height - 1)
+
+    return image, label
+
+
+class RandomCrop(object):
+    def __init__(self, size: Tuple[int, int]) -> None:
+        self.size = size
+        assert len(self.size) == 2, f"size should be a tuple (h, w), got {self.size}."
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        crop_height, crop_width = self.size
+        image_height, image_width = image.shape[-2:]
+        assert crop_height <= image_height and crop_width <= image_width, \
+            f"crop size should be no larger than image size, got crop size {self.size} and image size {image.shape}."
+        
+        top = torch.randint(0, image_height - crop_height + 1, (1,)).item()
+        left = torch.randint(0, image_width - crop_width + 1, (1,)).item()
+        return _crop(image, label, top, left, crop_height, crop_width)
+
+
+class Resize(object):
+    def __init__(self, size: Tuple[int, int]) -> None:
+        self.size = size
+        assert len(self.size) == 2, f"size should be a tuple (h, w), got {self.size}."
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        return _resize(image, label, self.size[0], self.size[1])
+
+
+class Resize2Multiple(object):
+    """
+    Resize the image so that it satisfies:
+        img_h = window_h + stride_h * n_h
+        img_w = window_w + stride_w * n_w
+    """
+    def __init__(
+        self,
+        window_size: Tuple[int, int],
+        stride: Tuple[int, int],
+    ) -> None:
+        window_size = (int(window_size), int(window_size)) if isinstance(window_size, (int, float)) else window_size
+        window_size = tuple(window_size)
+        stride = (int(stride), int(stride)) if isinstance(stride, (int, float)) else stride
+        stride = tuple(stride)
+        assert len(window_size) == 2, f"window_size should be a tuple (h, w), got {window_size}."
+        assert len(stride) == 2, f"stride should be a tuple (h, w), got {stride}."
+        assert all(s > 0 for s in window_size), f"window_size should be positive, got {window_size}."
+        assert all(s > 0 for s in stride), f"stride should be positive, got {stride}."
+        assert stride[0] <= window_size[0] and stride[1] <= window_size[1], f"stride should be no larger than window_size, got {stride} and {window_size}."
+        self.window_size = window_size
+        self.stride = stride
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        image_height, image_width = image.shape[-2:]
+        window_height, window_width = self.window_size
+        stride_height, stride_width = self.stride
+        new_height = int(max(round((image_height - window_height) / stride_height), 0) * stride_height + window_height)
+        new_width = int(max(round((image_width - window_width) / stride_width), 0) * stride_width + window_width)
+
+        if new_height == image_height and new_width == image_width:
+            return image, label
+        else:
+            return _resize(image, label, new_height, new_width)
+
+
+class ZeroPad2Multiple(object):
+    def __init__(
+        self,
+        window_size: Tuple[int, int],
+        stride: Tuple[int, int],
+    ) -> None:
+        window_size = (int(window_size), int(window_size)) if isinstance(window_size, (int, float)) else window_size
+        window_size = tuple(window_size)
+        stride = (int(stride), int(stride)) if isinstance(stride, (int, float)) else stride
+        stride = tuple(stride)
+        assert len(window_size) == 2, f"window_size should be a tuple (h, w), got {window_size}."
+        assert len(stride) == 2, f"stride should be a tuple (h, w), got {stride}."
+        assert all(s > 0 for s in window_size), f"window_size should be positive, got {window_size}."
+        assert all(s > 0 for s in stride), f"stride should be positive, got {stride}."
+        assert stride[0] <= window_size[0] and stride[1] <= window_size[1], f"stride should be no larger than window_size, got {stride} and {window_size}."
+        self.window_size = window_size
+        self.stride = stride
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        image_height, image_width = image.shape[-2:]
+        window_height, window_width = self.window_size
+        stride_height, stride_width = self.stride
+        new_height = int(max(np.ceil((image_height - window_height) / stride_height), 0) * stride_height + window_height)
+        new_width = int(max(np.ceil((image_width - window_width) / stride_width), 0) * stride_width + window_width)
+
+        if new_height == image_height and new_width == image_width:
+            return image, label
+        else:
+            assert new_height >= image_height and new_width >= image_width, f"new size should be no less than the original size, got {new_height} and {new_width}."
+            pad_height, pad_width = new_height - image_height, new_width - image_width
+            return TF.pad(image, (0, 0, pad_width, pad_height), fill=0), label  # only pad the right and bottom sides so that the label coordinates are not affected
+
+
+class RandomResizedCrop(object):
+    def __init__(
+        self,
+        size: Tuple[int, int],
+        scale: Tuple[float, float] = (0.75, 1.25),
+    ) -> None:
+        """
+        Randomly crop an image and resize it to a given size. The aspect ratio is preserved during this process.
+        """
+        self.size = size
+        self.scale = scale
+        assert len(self.size) == 2, f"size should be a tuple (h, w), got {self.size}."
+        assert 0 < self.scale[0] <= self.scale[1], f"scale should satisfy 0 < scale[0] <= scale[1], got {self.scale}."
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        out_height, out_width = self.size
+        # out_ratio = out_width / out_height
+
+        scale = torch.empty(1).uniform_(self.scale[0], self.scale[1]).item()  # if scale < 1, then the image will be zoomed in, otherwise zoomed out
+        in_height, in_width = image.shape[-2:]
+
+        # if in_width / in_height < out_ratio:  # Image is too tall
+        #     crop_width = int(in_width * scale)
+        #     crop_height = int(crop_width / out_ratio)
+        # else:  # Image is too wide
+        #     crop_height = int(in_height * scale)
+        #     crop_width = int(crop_height * out_ratio)
+
+        crop_height, crop_width = int(out_height * scale), int(out_width * scale)
+
+        if crop_height <= in_height and crop_width <= in_width:  # directly crop and resize the image
+            top = torch.randint(0, in_height - crop_height + 1, (1,)).item()
+            left = torch.randint(0, in_width - crop_width + 1, (1,)).item()
+
+        else:  # resize the image and then crop
+            ratio = max(crop_height / in_height, crop_width / in_width)  # keep the aspect ratio
+            resize_height, resize_width = int(in_height * ratio) + 1, int(in_width * ratio) + 1  # add 1 to make sure the resized image is no less than the crop size
+            image, label = _resize(image, label, resize_height, resize_width)
+            
+            top = torch.randint(0, resize_height - crop_height + 1, (1,)).item()
+            left = torch.randint(0, resize_width - crop_width + 1, (1,)).item()
+
+        image, label = _crop(image, label, top, left, crop_height, crop_width)
+        return _resize(image, label, out_height, out_width)
+        
+
+class RandomHorizontalFlip(object):
+    def __init__(self, p: float = 0.5) -> None:
+        self.p = p
+        assert 0 <= self.p <= 1, f"p should be in range [0, 1], got {self.p}."
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        if torch.rand(1) < self.p:
+            image = TF.hflip(image)
+
+            if len(label) > 0:
+                label[:, 0] = image.shape[-1] - 1 - label[:, 0]  # if width is 256, then 0 -> 255, 1 -> 254, 2 -> 253, etc.
+                label[:, 0] = label[:, 0].clamp(min=0, max=image.shape[-1] - 1)
+
+        return image, label
+    
+
+class ColorJitter(object):
+    def __init__(
+        self,
+        brightness: Union[float, Tuple[float, float]] = 0.4,
+        contrast: Union[float, Tuple[float, float]] = 0.4,
+        saturation: Union[float, Tuple[float, float]] = 0.4,
+        hue: Union[float, Tuple[float, float]] = 0.2,
+    ) -> None:
+        self.color_jitter = _ColorJitter(brightness=brightness, contrast=contrast, saturation=saturation, hue=hue)
+    
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        return self.color_jitter(image), label
+    
+
+class RandomGrayscale(object):
+    def __init__(self, p: float = 0.1) -> None:
+        self.p = p
+        assert 0 <= self.p <= 1, f"p should be in range [0, 1], got {self.p}."
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        if torch.rand(1) < self.p:
+            image = TF.rgb_to_grayscale(image, num_output_channels=3)
+
+        return image, label
+    
+
+class GaussianBlur(object):
+    def __init__(self, kernel_size: int, sigma: Optional[float] = None) -> None:
+        self.kernel_size = kernel_size
+        self.sigma = sigma
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        return TF.gaussian_blur(image, self.kernel_size, self.sigma), label
+
+
+class RandomApply(object):
+    def __init__(self, transforms: Tuple[Callable, ...], p: Union[float, Tuple[float, ...]] = 0.5) -> None:
+        self.transforms = transforms
+        p = [p] * len(transforms) if isinstance(p, float) else p
+        assert all(0 <= p_ <= 1 for p_ in p), f"p should be in range [0, 1], got {p}."
+        assert len(p) == len(transforms), f"p should be a float or a tuple of floats with the same length as transforms, got {p}."
+        self.p = p
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        for transform, p in zip(self.transforms, self.p):
+            if torch.rand(1) < p:
+                image, label = transform(image, label)
+
+        return image, label
+
+
+class PepperSaltNoise(object):
+    def __init__(self, saltiness: float = 0.001, spiciness: float = 0.001) -> None:
+        self.saltiness = saltiness
+        self.spiciness = spiciness
+        assert 0 <= self.saltiness <= 1, f"saltiness should be in range [0, 1], got {self.saltiness}."
+        assert 0 <= self.spiciness <= 1, f"spiciness should be in range [0, 1], got {self.spiciness}."
+
+    def __call__(self, image: Tensor, label: Tensor) -> Tuple[Tensor, Tensor]:
+        noise = torch.rand_like(image)
+        image = torch.where(noise < self.saltiness, 1., image)  # Salt
+        image = torch.where(noise > 1 - self.spiciness, 0., image)    # Pepper
+        return image, label

datasets/utils.py

@@ -0,0 +1,63 @@
+import torch
+from torch import Tensor
+from scipy.ndimage import gaussian_filter
+from typing import Optional, List, Tuple
+
+
+def get_id(x: str) -> int:
+    return int(x.split(".")[0])
+
+
+def generate_density_map(label: Tensor, height: int, width: int, sigma: Optional[float] = None) -> Tensor:
+    """
+    Generate the density map based on the dot annotations provided by the label.
+    """
+    density_map = torch.zeros((1, height, width), dtype=torch.float32)
+
+    if len(label) > 0:
+        assert len(label.shape) == 2 and label.shape[1] == 2, f"label should be a Nx2 tensor, got {label.shape}."
+        label_ = label.long()
+        label_[:, 0] = label_[:, 0].clamp(min=0, max=width - 1)
+        label_[:, 1] = label_[:, 1].clamp(min=0, max=height - 1)
+        density_map[0, label_[:, 1], label_[:, 0]] = 1.0
+
+    if sigma is not None:
+        assert sigma > 0, f"sigma should be positive if not None, got {sigma}."
+        density_map = torch.from_numpy(gaussian_filter(density_map, sigma=sigma))
+
+    return density_map
+
+
+def collate_fn(batch: List[Tensor]) -> Tuple[Tensor, List[Tensor], Tensor]:
+    batch = list(zip(*batch))
+    images = batch[0]
+    assert len(images[0].shape) == 4, f"images should be a 4D tensor, got {images[0].shape}."
+    if len(batch) == 4:  # image, label, density_map, image_name
+        images = torch.cat(images, 0)
+        points = batch[1]  # list of lists of tensors, flatten it
+        points = [p for points_ in points for p in points_]
+        densities = torch.cat(batch[2], 0)
+        image_names = batch[3]  # list of lists of strings, flatten it
+        image_names = [name for names_ in image_names for name in names_]
+
+        return images, points, densities, image_names
+
+    elif len(batch) == 3:  # image, label, density_map
+        images = torch.cat(images, 0)
+        points = batch[1]
+        points = [p for points_ in points for p in points_]
+        densities = torch.cat(batch[2], 0)
+
+        return images, points, densities
+    
+    elif len(batch) == 2:  # image, image_name. NWPU test dataset
+        images = torch.cat(images, 0)
+        image_names = batch[1]
+        image_names = [name for names_ in image_names for name in names_]
+
+        return images, image_names
+
+    else:
+        images = torch.cat(images, 0)
+
+        return images

eval.py

@@ -0,0 +1,41 @@
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+import numpy as np
+from typing import Dict, Optional
+from tqdm import tqdm
+
+from utils import calculate_errors, sliding_window_predict
+
+
+def evaluate(
+    model: nn.Module,
+    data_loader: DataLoader,
+    device: torch.device,
+    sliding_window: bool = False,
+    window_size: Optional[int] = None,
+    stride: Optional[int] = None,
+) -> Dict[str, float]:
+    model.eval()
+    pred_counts, target_counts = [], []
+    if sliding_window:
+        assert window_size is not None, f"Window size must be provided when sliding_window is True, but got {window_size}"
+        assert stride is not None, f"Stride must be provided when sliding_window is True, but got {stride}"
+
+    for image, target_points, _ in tqdm(data_loader):
+        image = image.to(device)
+        target_counts.append([len(p) for p in target_points])
+
+        with torch.set_grad_enabled(False):
+            # if sliding_window:
+            #     pred_density = sliding_window_predict(model, image, window_size, stride)
+            # else:
+            #     pred_density = model(image)
+            pred_density = model(image)
+
+            pred_counts.append(pred_density.sum(dim=(1, 2, 3)).cpu().numpy().tolist())
+
+    pred_counts = np.array([item for sublist in pred_counts for item in sublist])
+    target_counts = np.array([item for sublist in target_counts for item in sublist])
+    assert len(pred_counts) == len(target_counts), f"Length of predictions and ground truths should be equal, but got {len(pred_counts)} and {len(target_counts)}"
+    return calculate_errors(pred_counts, target_counts)

losses/__init__.py

@@ -0,0 +1,7 @@
+from .dm_loss import DMLoss
+from .dace_loss import DACELoss
+
+__all__ = [
+    "DMLoss",
+    "DACELoss",
+]

losses/bregman_pytorch.py

@@ -0,0 +1,144 @@
+# Code modified from https://github.com/cvlab-stonybrook/DM-Count/blob/master/losses/bregman_pytorch.py
+import torch
+from torch import Tensor
+from torch.cuda.amp import autocast
+from typing import Union, Tuple, Dict
+
+M_EPS = 1e-16
+
+
+@autocast(enabled=True, dtype=torch.float32)  # avoid numerical instability
+def sinkhorn(
+    a: Tensor,
+    b: Tensor,
+    C: Tensor,
+    reg: float = 1e-1,
+    maxIter: int = 1000,
+    stopThr: float = 1e-9,
+    verbose: bool = False,
+    log: bool = True,
+    eval_freq: int = 10,
+    print_freq: int = 200,
+) -> Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]:
+    """
+    Solve the entropic regularization optimal transport
+    The input should be PyTorch tensors
+    The function solves the following optimization problem:
+
+    .. math::
+        \gamma = arg\min_\gamma <\gamma,C>_F + reg\cdot\Omega(\gamma)
+        s.t.    \gamma 1 = a
+                \gamma^T 1= b
+                \gamma\geq 0
+    where :
+    - C is the (ns,nt) metric cost matrix
+    - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
+    - a and b are target and source measures (sum to 1)
+    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [1].
+
+    Parameters
+    ----------
+    a : torch.tensor (na,)
+        samples measure in the target domain
+    b : torch.tensor (nb,)
+        samples in the source domain
+    C : torch.tensor (na,nb)
+        loss matrix
+    reg : float
+        Regularization term > 0
+    maxIter : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error ( > 0 )
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+    Returns
+    -------
+    gamma : (na x nb) torch.tensor
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    References
+    ----------
+    [1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
+    See Also
+    --------
+    """
+
+    device = a.device
+    na, nb = C.shape
+
+    # a = a.view(-1, 1)
+    # b = b.view(-1, 1)
+
+    assert na >= 1 and nb >= 1, f"C needs to be 2d. Found C.shape = {C.shape}"
+    assert na == a.shape[0] and nb == b.shape[0], f"Shape of a ({a.shape}) or b ({b.shape}) does not match that of C ({C.shape})"
+    assert reg > 0, f"reg should be greater than 0. Found reg = {reg}"
+    assert a.min() >= 0. and b.min() >= 0., f"Elements in a and b should be nonnegative. Found a.min() = {a.min()}, b.min() = {b.min()}"
+
+    if log:
+        log = {"err": []}
+
+    u = torch.ones((na), dtype=a.dtype).to(device) / na
+    v = torch.ones((nb), dtype=b.dtype).to(device) / nb
+
+    K = torch.empty(C.shape, dtype=C.dtype).to(device)
+    torch.div(C, -reg, out=K)
+    torch.exp(K, out=K)
+
+    b_hat = torch.empty(b.shape, dtype=C.dtype).to(device)
+
+    it = 1
+    err = 1
+
+    # allocate memory beforehand
+    KTu = torch.empty(v.shape, dtype=v.dtype).to(device)
+    Kv = torch.empty(u.shape, dtype=u.dtype).to(device)
+
+    while (err > stopThr and it <= maxIter):
+        upre, vpre = u, v
+        # torch.matmul(u, K, out=KTu)
+        KTu = torch.matmul(u.view(1, -1), K).view(-1)
+        v = torch.div(b, KTu + M_EPS)
+        # torch.matmul(K, v, out=Kv)
+        Kv = torch.matmul(K, v.view(-1, 1)).view(-1)
+        u = torch.div(a, Kv + M_EPS)
+
+        if torch.any(torch.isnan(u)) or torch.any(torch.isnan(v)) or \
+                torch.any(torch.isinf(u)) or torch.any(torch.isinf(v)):
+            print("Warning: numerical errors at iteration", it)
+            u, v = upre, vpre
+            break
+
+        if log and it % eval_freq == 0:
+            # we can speed up the process by checking for the error only all
+            # the eval_freq iterations
+            # below is equivalent to:
+            # b_hat = torch.sum(u.reshape(-1, 1) * K * v.reshape(1, -1), 0)
+            # but with more memory efficient
+            b_hat = (torch.matmul(u.view(1, -1), K) * v.view(1, -1)).view(-1)
+            err = (b - b_hat).pow(2).sum().item()
+            # err = (b - b_hat).abs().sum().item()
+            log["err"].append(err)
+
+        if verbose and it % print_freq == 0:
+            print("iteration {:5d}, constraint error {:5e}".format(it, err))
+
+        it += 1
+
+    if log:
+        log["u"] = u
+        log["v"] = v
+        log["alpha"] = reg * torch.log(u + M_EPS)
+        log["beta"] = reg * torch.log(v + M_EPS)
+
+    # transport plan
+    P = u.reshape(-1, 1) * K * v.reshape(1, -1)
+    if log:
+        return P, log
+    else:
+        return P

losses/dace_loss.py

@@ -0,0 +1,70 @@
+import torch
+from torch import nn, Tensor
+from typing import Any, List, Tuple, Dict
+
+from .dm_loss import DMLoss
+from .utils import _reshape_density
+
+
+class DACELoss(nn.Module):
+    def __init__(
+        self,
+        bins: List[Tuple[float, float]],
+        reduction: int,
+        weight_count_loss: float = 1.0,
+        count_loss: str = "mae",
+        **kwargs: Any
+    ) -> None:
+        super().__init__()
+        assert len(bins) > 0, f"Expected at least one bin, got {bins}"
+        assert all([len(b) == 2 for b in bins]), f"Expected all bins to be of length 2, got {bins}"
+        assert all([b[0] <= b[1] for b in bins]), f"Expected all bins to be in increasing order, got {bins}"
+        self.bins = bins
+        self.reduction = reduction
+        self.cross_entropy_fn = nn.CrossEntropyLoss(reduction="none")
+
+        count_loss = count_loss.lower()
+        assert count_loss in ["mae", "mse", "dmcount"], f"Expected count_loss to be one of ['mae', 'mse', 'dmcount'], got {count_loss}"
+        self.count_loss = count_loss
+        if self.count_loss == "mae":
+            self.use_dm_loss = False
+            self.count_loss_fn = nn.L1Loss(reduction="none")
+        elif self.count_loss == "mse":
+            self.use_dm_loss = False
+            self.count_loss_fn = nn.MSELoss(reduction="none")
+        else:
+            self.use_dm_loss = True
+            assert "input_size" in kwargs, f"Expected input_size to be in kwargs when count_loss='dmcount', got {kwargs}"
+            self.count_loss_fn = DMLoss(reduction=reduction, **kwargs)
+
+        self.weight_count_loss = weight_count_loss
+
+    def _bin_count(self, density_map: Tensor) -> Tensor:
+        class_map = torch.zeros_like(density_map, dtype=torch.long)
+        for idx, (low, high) in enumerate(self.bins):
+            mask = (density_map >= low) & (density_map <= high)
+            class_map[mask] = idx
+        return class_map.squeeze(1)  # remove channel dimension
+
+    def forward(self, pred_class: Tensor, pred_density: Tensor, target_density: Tensor, target_points: List[Tensor]) -> Tuple[Tensor, Dict[str, Tensor]]:
+        target_density = _reshape_density(target_density, reduction=self.reduction) if target_density.shape[-2:] != pred_density.shape[-2:] else target_density
+        assert pred_density.shape == target_density.shape, f"Expected pred_density and target_density to have the same shape, got {pred_density.shape} and {target_density.shape}"
+
+        target_class = self._bin_count(target_density)
+
+        cross_entropy_loss = self.cross_entropy_fn(pred_class, target_class).sum(dim=(-1, -2)).mean()
+
+        if self.use_dm_loss:
+            count_loss, loss_info = self.count_loss_fn(pred_density, target_density, target_points)
+            loss_info["ce_loss"] = cross_entropy_loss.detach()
+        else:
+            count_loss = self.count_loss_fn(pred_density, target_density).sum(dim=(-1, -2, -3)).mean()
+            loss_info = {
+                "ce_loss": cross_entropy_loss.detach(),
+                f"{self.count_loss}_loss": count_loss.detach(),
+            }
+
+        loss = cross_entropy_loss + self.weight_count_loss * count_loss
+        loss_info["loss"] = loss.detach()
+
+        return loss, loss_info

losses/dm_loss.py

@@ -0,0 +1,124 @@
+import torch
+from torch import nn, Tensor
+from torch.cuda.amp import autocast
+from typing import List, Any, Tuple, Dict
+
+from .bregman_pytorch import sinkhorn
+from .utils import _reshape_density
+
+EPS = 1e-8
+
+
+class OTLoss(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        reduction: int,
+        norm_cood: bool,
+        num_of_iter_in_ot: int = 100,
+        reg: float = 10.0
+    ) -> None:
+        super().__init__()
+        assert input_size % reduction == 0
+
+        self.input_size = input_size
+        self.reduction = reduction
+        self.norm_cood = norm_cood
+        self.num_of_iter_in_ot = num_of_iter_in_ot
+        self.reg = reg
+
+        # coordinate is same to image space, set to constant since crop size is same
+        self.cood = torch.arange(0, input_size, step=reduction, dtype=torch.float32) + reduction / 2
+        self.density_size = self.cood.size(0)
+        self.cood.unsqueeze_(0) # [1, #cood]
+        self.cood = self.cood / input_size * 2 - 1 if self.norm_cood else self.cood
+        self.output_size = self.cood.size(1)
+
+    @autocast(enabled=True, dtype=torch.float32)  # avoid numerical instability
+    def forward(self, pred_density: Tensor, normed_pred_density: Tensor, target_points: List[Tensor]) -> Tuple[Tensor, float, Tensor]:
+        batch_size = normed_pred_density.size(0)
+        assert len(target_points) == batch_size, f"Expected target_points to have length {batch_size}, but got {len(target_points)}"
+        assert self.output_size == normed_pred_density.size(2)
+        device = pred_density.device
+
+        loss = torch.zeros([1]).to(device)
+        ot_obj_values = torch.zeros([1]).to(device)
+        wd = 0 # Wasserstein distance
+        cood = self.cood.to(device)
+        for idx, points in enumerate(target_points):
+            if len(points) > 0:
+                # compute l2 square distance, it should be source target distance. [#gt, #cood * #cood]
+                points = points / self.input_size * 2 - 1 if self.norm_cood else points
+                x = points[:, 0].unsqueeze_(1)  # [#gt, 1]
+                y = points[:, 1].unsqueeze_(1)
+                x_dist = -2 * torch.matmul(x, cood) + x * x + cood * cood # [#gt, #cood]
+                y_dist = -2 * torch.matmul(y, cood) + y * y + cood * cood
+                y_dist.unsqueeze_(2)
+                x_dist.unsqueeze_(1)
+                dist = y_dist + x_dist
+                dist = dist.view((dist.size(0), -1)) # size of [#gt, #cood * #cood]
+
+                source_prob = normed_pred_density[idx][0].view([-1]).detach()
+                target_prob = (torch.ones([len(points)]) / len(points)).to(device)
+                # use sinkhorn to solve OT, compute optimal beta.
+                P, log = sinkhorn(target_prob, source_prob, dist, self.reg, maxIter=self.num_of_iter_in_ot, log=True)
+                beta = log["beta"] # size is the same as source_prob: [#cood * #cood]
+                ot_obj_values += torch.sum(normed_pred_density[idx] * beta.view([1, self.output_size, self.output_size]))
+                # compute the gradient of OT loss to predicted density (pred_density).
+                # im_grad = beta / source_count - < beta, source_density> / (source_count)^2
+                source_density = pred_density[idx][0].view([-1]).detach()
+                source_count = source_density.sum()
+                gradient_1 = (source_count) / (source_count * source_count+ EPS) * beta # size of [#cood * #cood]
+                gradient_2 = (source_density * beta).sum() / (source_count * source_count + EPS) # size of 1
+                gradient = gradient_1 - gradient_2
+                gradient = gradient.detach().view([1, self.output_size, self.output_size])
+                # Define loss = <im_grad, predicted density>. The gradient of loss w.r.t predicted density is im_grad.
+                loss += torch.sum(pred_density[idx] * gradient)
+                wd += torch.sum(dist * P).item()
+
+        return loss, wd, ot_obj_values
+
+
+class DMLoss(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        reduction: int,
+        norm_cood: bool = False,
+        weight_ot: float = 0.1,
+        weight_tv: float = 0.01,
+        **kwargs: Any
+    ) -> None:
+        super().__init__()
+        self.ot_loss = OTLoss(input_size, reduction, norm_cood, **kwargs)
+        self.tv_loss = nn.L1Loss(reduction="none")
+        self.count_loss = nn.L1Loss(reduction="mean")
+        self.weight_ot = weight_ot
+        self.weight_tv = weight_tv
+
+    @autocast(enabled=True, dtype=torch.float32)  # avoid numerical instability
+    def forward(self, pred_density: Tensor, target_density: Tensor, target_points: List[Tensor]) -> Tuple[Tensor, Dict[str, Tensor]]:
+        target_density = _reshape_density(target_density, reduction=self.ot_loss.reduction) if target_density.shape[-2:] != pred_density.shape[-2:] else target_density
+        assert pred_density.shape == target_density.shape, f"Expected pred_density and target_density to have the same shape, got {pred_density.shape} and {target_density.shape}"
+
+        pred_count = pred_density.view(pred_density.shape[0], -1).sum(dim=1)
+        normed_pred_density = pred_density / (pred_count.view(-1, 1, 1, 1) + EPS)
+        target_count = torch.tensor([len(p) for p in target_points], dtype=torch.float32).to(target_density.device)
+        normed_target_density = target_density / (target_count.view(-1, 1, 1, 1) + EPS)
+
+        ot_loss, _, _ = self.ot_loss(pred_density, normed_pred_density, target_points)
+
+        tv_loss = (self.tv_loss(normed_pred_density, normed_target_density).sum(dim=(1, 2, 3)) * target_count).mean()
+
+        count_loss = self.count_loss(pred_count, target_count)
+
+        loss = ot_loss * self.weight_ot + tv_loss * self.weight_tv + count_loss
+
+        loss_info = {
+            "loss": loss.detach(),
+            "ot_loss": ot_loss.detach(),
+            "tv_loss": tv_loss.detach(),
+            "count_loss": count_loss.detach(),
+        }
+
+        return loss, loss_info

losses/utils.py

@@ -0,0 +1,9 @@
+from torch import Tensor
+
+
+def _reshape_density(density: Tensor, reduction: int) -> Tensor:
+    assert len(density.shape) == 4, f"Expected 4D (B, 1, H, W) tensor, got {density.shape}"
+    assert density.shape[1] == 1, f"Expected 1 channel, got {density.shape[1]}"
+    assert density.shape[2] % reduction == 0, f"Expected height to be divisible by {reduction}, got {density.shape[2]}"
+    assert density.shape[3] % reduction == 0, f"Expected width to be divisible by {reduction}, got {density.shape[3]}"
+    return density.reshape(density.shape[0], 1, density.shape[2] // reduction, reduction, density.shape[3] // reduction, reduction).sum(dim=(-1, -3))

models/__init__.py

@@ -0,0 +1,49 @@
+from typing import List, Tuple, Optional, Any, Union
+
+from .model import _classifier, _regressor, Classifier, Regressor
+from .clip import _clip_ebc, CLIP_EBC
+
+
+clip_names = ["resnet50", "resnet50x4", "resnet50x16", "resnet50x64", "resnet101", "vit_b_16", "vit_b_32", "vit_l_14"]
+
+
+def get_model(
+    backbone: str,
+    input_size: int,
+    reduction: int,
+    bins: Optional[List[Tuple[float, float]]] = None,
+    anchor_points: Optional[List[float]] = None,
+    **kwargs: Any,
+) -> Union[Regressor, Classifier, CLIP_EBC]:
+    backbone = backbone.lower()
+    if "clip" in backbone:
+        backbone = backbone[5:]
+        assert backbone in clip_names, f"Expected backbone to be in {clip_names}, got {backbone}"
+        return _clip_ebc(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+            bins=bins,
+            anchor_points=anchor_points,
+            **kwargs
+        )
+    elif bins is None and anchor_points is None:
+        return _regressor(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+        )
+    else:
+        assert bins is not None and anchor_points is not None, f"Expected bins and anchor_points to be both None or not None, got {bins} and {anchor_points}"
+        return _classifier(
+            backbone=backbone,
+            input_size=input_size,
+            reduction=reduction,
+            bins=bins,
+            anchor_points=anchor_points,
+        )
+
+
+__all__ = [
+    "get_model",
+]

models/clip/__init__.py

@@ -0,0 +1,7 @@
+from .model import CLIP_EBC, _clip_ebc
+
+
+__all__ = [
+    "CLIP_EBC",
+    "_clip_ebc",
+]

models/clip/_clip/__init__.py

@@ -0,0 +1,273 @@
+import torch
+import os
+from typing import Tuple, Optional, Any, Union
+import json
+
+from .utils import tokenize, transform
+from .prepare import prepare
+from .text_encoder import CLIPTextEncoder
+from .image_encoder import ModifiedResNet, VisionTransformer
+from .model import CLIP
+
+
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+
+clip_model_names = [
+    "clip_resnet50",
+    "clip_resnet101",
+    "clip_resnet50x4",
+    "clip_resnet50x16",
+    "clip_resnet50x64",
+    "clip_vit_b_32",
+    "clip_vit_b_16",
+    "clip_vit_l_14",
+    "clip_vit_l_14_336px",
+]
+
+clip_image_encoder_names = [f"clip_image_encoder_{name[5:]}" for name in clip_model_names]
+clip_text_encoder_names = [f"clip_text_encoder_{name[5:]}" for name in clip_model_names]
+
+
+for name in clip_model_names + clip_image_encoder_names + clip_text_encoder_names:
+    model_weights_path = os.path.join(curr_dir, "weights", f"{name}.pth")
+    model_config_path = os.path.join(curr_dir, "configs", f"{name}.json")
+    if not os.path.exists(os.path.join(curr_dir, "weights", f"{name}.pth")) or not os.path.exists(os.path.join(curr_dir, "configs", f"{name}.json")):
+        prepare()
+        break
+
+
+for name in clip_model_names + clip_image_encoder_names + clip_text_encoder_names:
+    assert os.path.exists(os.path.join(curr_dir, "weights", f"{name}.pth")), f"Missing {name}.pth in weights folder. Please run models/clip/prepare.py to download the weights."
+    assert os.path.exists(os.path.join(curr_dir, "configs", f"{name}.json")), f"Missing {name}.json in configs folder. Please run models/clip/prepare.py to download the configs."
+
+
+def _clip(name: str, input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    with open(os.path.join(curr_dir, "configs", f"clip_{name}.json"), "r") as f:
+        config = json.load(f)
+
+    model = CLIP(
+        embed_dim=config["embed_dim"],
+        # vision
+        image_resolution=config["image_resolution"],
+        vision_layers=config["vision_layers"],
+        vision_width=config["vision_width"],
+        vision_patch_size=config["vision_patch_size"],
+        # text
+        context_length=config["context_length"],
+        vocab_size=config["vocab_size"],
+        transformer_width=config["transformer_width"],
+        transformer_heads=config["transformer_heads"],
+        transformer_layers=config["transformer_layers"]
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_{name}.pth"), map_location="cpu")
+    model.load_state_dict(state_dict, strict=True)
+
+    if input_size is not None:
+        input_size = (input_size, input_size) if isinstance(input_size, int) else input_size
+        if name.startswith("vit"):
+            model.visual.adjust_pos_embed(*input_size)
+
+    return model
+
+
+def _resnet(
+    name: str,
+    reduction: int = 32,
+    features_only: bool = False,
+    out_indices: Optional[Tuple[int, ...]] = None,
+    **kwargs: Any
+) -> ModifiedResNet:
+    with open(os.path.join(curr_dir, "configs", f"clip_image_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = ModifiedResNet(
+        layers=config["vision_layers"],
+        output_dim=config["embed_dim"],
+        input_resolution=config["image_resolution"],
+        width=config["vision_width"],
+        heads=config["vision_heads"],
+        features_only=features_only,
+        out_indices=out_indices,
+        reduction=reduction
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_image_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    return model
+
+
+def _vit(name: str, features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    with open(os.path.join(curr_dir, "configs", f"clip_image_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = VisionTransformer(
+        input_resolution=config["image_resolution"],
+        patch_size=config["vision_patch_size"],
+        output_dim=config["embed_dim"],
+        width=config["vision_width"],
+        layers=config["vision_layers"],
+        heads=config["vision_heads"],
+        features_only=features_only
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_image_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    if input_size is not None:
+        input_size = (input_size, input_size) if isinstance(input_size, int) else input_size
+        model.adjust_pos_embed(*input_size)
+    return model
+
+
+def _text_encoder(name: str) -> CLIPTextEncoder:
+    with open(os.path.join(curr_dir, "configs", f"clip_text_encoder_{name}.json"), "r") as f:
+        config = json.load(f)
+    model = CLIPTextEncoder(
+        embed_dim=config["embed_dim"],
+        context_length=config["context_length"],
+        vocab_size=config["vocab_size"],
+        transformer_width=config["transformer_width"],
+        transformer_heads=config["transformer_heads"],
+        transformer_layers=config["transformer_layers"]
+    )
+    state_dict = torch.load(os.path.join(curr_dir, "weights", f"clip_text_encoder_{name}.pth"), map_location="cpu")
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+        print(f"Unexpected keys: {unexpected_keys}")
+    else:
+        print(f"All keys matched successfully.")
+
+    return model
+
+
+
+# CLIP models
+def resnet50_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50", input_size)
+
+def resnet101_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet101", input_size)
+
+def resnet50x4_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x4", input_size)
+
+def resnet50x16_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x16", input_size)
+
+def resnet50x64_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("resnet50x64", input_size)
+
+def vit_b_32_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_b_32", input_size)
+
+def vit_b_16_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_b_16", input_size)
+
+def vit_l_14_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_l_14", input_size)
+
+def vit_l_14_336px_clip(input_size: Optional[Union[int, Tuple[int, int]]] = None) -> CLIP:
+    return _clip("vit_l_14_336px", input_size)
+
+
+# CLIP image encoders
+def resnet50_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet101_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet101", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x4_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x4", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x16_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x16", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def resnet50x64_img(features_only: bool = False, out_indices: Optional[Tuple[int, ...]] = None, **kwargs: Any) -> ModifiedResNet:
+    return _resnet("resnet50x64", features_only=features_only, out_indices=out_indices, **kwargs)
+
+def vit_b_32_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_b_32", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_b_16_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_b_16", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_l_14_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_l_14", features_only=features_only, input_size=input_size, **kwargs)
+
+def vit_l_14_336px_img(features_only: bool = False, input_size: Optional[Union[int, Tuple[int, int]]] = None, **kwargs: Any) -> VisionTransformer:
+    return _vit("vit_l_14_336px", features_only=features_only, input_size=input_size, **kwargs)
+
+
+# CLIP text encoders
+def resnet50_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50")
+
+def resnet101_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet101")
+
+def resnet50x4_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x4")
+
+def resnet50x16_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x16")
+
+def resnet50x64_txt() -> CLIPTextEncoder:
+    return _text_encoder("resnet50x64")
+
+def vit_b_32_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_b_32")
+
+def vit_b_16_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_b_16")
+
+def vit_l_14_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_l_14")
+
+def vit_l_14_336px_txt() -> CLIPTextEncoder:
+    return _text_encoder("vit_l_14_336px")
+
+
+__all__ = [
+    # utils
+    "tokenize",
+    "transform",
+    # clip models
+    "resnet50_clip",
+    "resnet101_clip",
+    "resnet50x4_clip",
+    "resnet50x16_clip",
+    "resnet50x64_clip",
+    "vit_b_32_clip",
+    "vit_b_16_clip",
+    "vit_l_14_clip",
+    "vit_l_14_336px_clip",
+    # clip image encoders
+    "resnet50_img",
+    "resnet101_img",
+    "resnet50x4_img",
+    "resnet50x16_img",
+    "resnet50x64_img",
+    "vit_b_32_img",
+    "vit_b_16_img",
+    "vit_l_14_img",
+    "vit_l_14_336px_img",
+    # clip text encoders
+    "resnet50_txt",
+    "resnet101_txt",
+    "resnet50x4_txt",
+    "resnet50x16_txt",
+    "resnet50x64_txt",
+    "vit_b_32_txt",
+    "vit_b_16_txt",
+    "vit_l_14_txt",
+    "vit_l_14_336px_txt",
+]

models/clip/_clip/blocks.py

@@ -0,0 +1,137 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from collections import OrderedDict
+from typing import Optional, Iterable
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: 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: Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: 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: 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: Tensor) -> 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: 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: Tensor):
+        return self.resblocks(x)
+
+
+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.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        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.relu3 = 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: Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(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.relu3(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 + 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.flatten(start_dim=2).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[:1], 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.squeeze(0)

models/clip/_clip/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

models/clip/_clip/image_encoder.py

@@ -0,0 +1,225 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Tuple, Union, Any, List, Iterable, Optional
+
+from .blocks import LayerNorm, Transformer, Bottleneck, AttentionPool2d
+
+
+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: Tuple[int, int, int, int],
+        output_dim: int,
+        input_resolution: int = 224,
+        width: int = 64,
+        heads: int = 8,
+        features_only: bool = False,
+        out_indices: Optional[Iterable[int]] = None,
+        reduction: int = 32,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        input_resolution = (input_resolution, input_resolution) if isinstance(input_resolution, int) else input_resolution
+        assert isinstance(input_resolution, tuple) and len(input_resolution) == 2, f"input_resolution should be a tuple of length 2, but got {input_resolution}"
+        self.input_resolution = input_resolution
+        self.downsampling_rate = 32  # the rate at which the input is downsampled by the network
+
+        # 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.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # 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=1 if reduction <= 16 else 2)
+
+        self.features_only = features_only
+        if features_only:
+            self.out_indices = out_indices if out_indices is not None else range(5)
+            self.out_indices = [idx + 5 if idx < 0 else idx for idx in self.out_indices]  # map negative indices to positive indices
+            self.out_indices = sorted(set(self.out_indices))  # remove duplicates and sort
+            assert min(self.out_indices) >= 0 and max(self.out_indices) <= 4, f"out_indices={self.out_indices} is invalid for a ResNet with 5 stages"
+            self.channels = width * 32  # the ResNet feature dimension
+        else:
+            self.out_indices = None
+            embed_dim = width * 32  # the ResNet feature dimension
+            self.attnpool = AttentionPool2d((input_resolution[0] // 32) * (input_resolution[1] // 32), embed_dim, heads, output_dim)
+            self.channels = output_dim
+
+        self.reduction = self.downsampling_rate // 2 if reduction <= 16 else self.downsampling_rate
+        self.clip_embed_dim = 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 _stem(self, x: Tensor) -> Tensor:
+        x = self.relu1(self.bn1(self.conv1(x)))
+        x = self.relu2(self.bn2(self.conv2(x)))
+        x = self.relu3(self.bn3(self.conv3(x)))
+        x = self.avgpool(x)
+        return x
+
+    def forward(self, x: Tensor) -> Union[Tensor, List[Tensor]]:
+        x = x.type(self.conv1.weight.dtype)
+        x = self._stem(x)
+
+        feats = [x] if self.features_only and 0 in self.out_indices else []
+
+        x = self.layer1(x)
+        if self.features_only and 1 in self.out_indices:
+            feats.append(x)
+
+        x = self.layer2(x)
+        if self.features_only and 2 in self.out_indices:
+            feats.append(x)
+
+        x = self.layer3(x)
+        if self.features_only and 3 in self.out_indices:
+            feats.append(x)
+
+        x = self.layer4(x)
+        if self.features_only and 4 in self.out_indices:
+            feats.append(x)
+
+        if self.features_only:
+            if len(self.out_indices) == 1:
+                return feats[0]
+            else:
+                return feats
+        else:
+            x = self.attnpool(x)
+            return x
+
+
+class VisionTransformer(nn.Module):
+    def __init__(
+        self,
+        input_resolution: Union[int, Tuple[int, int]],
+        patch_size: Union[int, Tuple[int, int]],
+        output_dim: int,
+        width: int,
+        layers: int,
+        heads: int,
+        features_only: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        input_resolution = (input_resolution, input_resolution) if isinstance(input_resolution, int) else input_resolution
+        patch_size = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size
+        assert isinstance(input_resolution, tuple) and len(input_resolution) == 2, f"input_resolution should be a tuple of length 2, but got {input_resolution}"
+        assert isinstance(patch_size, tuple) and len(patch_size) == 2, f"patch_size should be a tuple of length 2, but got {patch_size}"
+        assert patch_size[0] == patch_size[1], f"ViT only supports square patches, patch_size={patch_size} is invalid."
+        assert input_resolution[0] % patch_size[0] == 0 and input_resolution[1] % patch_size[1] == 0, f"input_resolution {input_resolution} should be divisible by patch_size {patch_size}"
+        self.input_resolution = input_resolution
+        self.patch_size = patch_size
+        self.downsampling_rate = patch_size[0]
+
+        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.num_patches_h = int(input_resolution[0] // patch_size[0])
+        self.num_patches_w = int(input_resolution[1] // patch_size[1])
+        self.positional_embedding = nn.Parameter(scale * torch.randn(self.num_patches_h * self.num_patches_w + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+        self.ln_post = LayerNorm(width)
+
+        self.features_only = features_only  # if True, return the final patches instead of the CLS token
+        if features_only:
+            self.channels = width
+        else:
+            self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+            self.channels = output_dim
+
+        self.reduction = patch_size[0]
+        self.clip_embed_dim = output_dim
+
+    def adjust_pos_embed(self, h: int, w: int) -> None:
+        """
+        Permanently adjust the size of the positional embedding matrix.
+
+        Args:
+            h: the height of the original input image.
+            w: the width of the original input image.
+        """
+        assert h % self.patch_size[0] == 0 and w % self.patch_size[1] == 0, f"input_resolution {h, w} should be divisible by patch_size {self.patch_size}"
+        if self.input_resolution[0] != h or self.input_resolution[1] != w:
+            new_num_patches_h = int(h // self.patch_size[0])
+            new_num_patches_w = int(w // self.patch_size[1])
+            positional_embedding = rearrange(self.positional_embedding[1:, :], "(h w) c -> c h w", h=self.num_patches_h, w=self.num_patches_w).unsqueeze(0)  # add batch dimension
+            positional_embedding = F.interpolate(positional_embedding, size=(new_num_patches_h, new_num_patches_w), mode="bicubic", ).squeeze(0)  # remove batch dimension
+            positional_embedding = rearrange(positional_embedding, "c h w -> (h w) c")
+            self.positional_embedding = nn.Parameter(torch.cat([self.positional_embedding[:1, :], positional_embedding], dim=0))
+            self.input_resolution = (h, w)
+            self.num_patches_h = new_num_patches_h
+            self.num_patches_w = new_num_patches_w
+
+    def _interpolate_pos_embed(self, h: int, w: int) -> Tensor:
+        """
+        Interpolate the positional embedding matrix to match the size of the input image.
+
+        Args:
+            h: the required number of patches along the height dimension.
+            w: the required number of patches along the width dimension.
+        """
+        if h == self.num_patches_h and w == self.num_patches_w:
+            return self.positional_embedding
+        else:
+            positional_embedding = rearrange(self.positional_embedding[1:, :], "(h w) c -> c h w", h=self.num_patches_h, w=self.num_patches_w).unsqueeze(0)  # add batch dimension
+            positional_embedding = F.interpolate(positional_embedding, size=(h, w), mode="bicubic").squeeze(0)  # remove batch dimension
+            positional_embedding = rearrange(positional_embedding, "c h w -> (h w) c")
+            positional_embedding = torch.cat([self.positional_embedding[:1, :], positional_embedding], dim=0)
+            return positional_embedding
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.conv1(x) # shape = [*, width, grid, grid]
+        num_patches_h, num_patches_w = x.shape[-2:]
+
+        positional_embedding = self._interpolate_pos_embed(num_patches_h, num_patches_w).to(x.dtype)
+        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)
+        x = x + positional_embedding
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND. N: batch size, L: sequence length, D: feature dimension
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_post(x)
+
+        if self.features_only:
+            x = x[:, 1:, :]  # remove the CLS token
+            x = rearrange(x, "n (h w) c -> n c h w", h=num_patches_h, w=num_patches_w)
+        else:
+            x = x[:, 0, :]
+            x = x @ self.proj
+        return x

models/clip/_clip/model.py

@@ -0,0 +1,214 @@
+import torch
+from torch import nn
+import numpy as np
+
+from typing import Tuple, Union
+
+from .image_encoder import ModifiedResNet, VisionTransformer
+from .text_encoder import LayerNorm, Transformer
+
+
+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
+        ) -> None:
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.image_resolution = image_resolution
+        self.vision_layers = vision_layers
+        self.vision_width = vision_width
+        self.vision_patch_size = vision_patch_size
+        self.context_length = context_length
+        self.vocab_size = vocab_size
+        self.transformer_width = transformer_width
+        self.transformer_heads = transformer_heads
+        self.transformer_layers = transformer_layers
+
+        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,
+                features_only=False,
+            )
+        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,
+                features_only=False,
+            )
+        self.vision_heads = vision_heads
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        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
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    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()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    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("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]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict, strict=False)
+    return model.eval()

models/clip/_clip/prepare.py

@@ -0,0 +1,95 @@
+# Prepare the models to speed up loading them later
+import torch
+from torch import nn, Tensor
+import os
+from tqdm import tqdm
+import json
+
+from .utils import load
+
+
+model_name_map = {
+    "RN50": "resnet50",
+    "RN101": "resnet101",
+    "RN50x4": "resnet50x4",
+    "RN50x16": "resnet50x16",
+    "RN50x64": "resnet50x64",
+    "ViT-B/32": "vit_b_32",
+    "ViT-B/16": "vit_b_16",
+    "ViT-L/14": "vit_l_14",
+    "ViT-L/14@336px": "vit_l_14_336px",
+}
+
+
+class CLIPTextEncoderTemp(nn.Module):
+    def __init__(
+        self,
+        clip: nn.Module,
+    ) -> None:
+        super().__init__()
+        self.context_length = clip.context_length
+        self.vocab_size = clip.vocab_size
+        self.dtype = clip.dtype
+        self.token_embedding = clip.token_embedding
+        self.positional_embedding = clip.positional_embedding
+        self.transformer = clip.transformer
+        self.ln_final = clip.ln_final
+        self.text_projection = clip.text_projection
+
+    def forward(self, text: Tensor) -> None:
+        pass
+
+
+def prepare() -> None:
+    print("Preparing CLIP models...")
+    curr_dir = os.path.dirname(os.path.abspath(__file__))
+    weight_dir = os.path.join(curr_dir, "weights")
+    config_dir = os.path.join(curr_dir, "configs")
+    os.makedirs(weight_dir, exist_ok=True)
+    os.makedirs(config_dir, exist_ok=True)
+    device = torch.device("cpu")
+
+    for model_name in tqdm(["RN50", "RN101", "RN50x4", "RN50x16", "RN50x64", "ViT-B/32", "ViT-B/16", "ViT-L/14", "ViT-L/14@336px"]):
+        model = load(model_name, device=device).to(device)
+        image_encoder = model.visual.to(device)
+        text_encoder = CLIPTextEncoderTemp(model).to(device)
+        torch.save(model.state_dict(), os.path.join(weight_dir, f"clip_{model_name_map[model_name]}.pth"))
+        torch.save(image_encoder.state_dict(), os.path.join(weight_dir, f"clip_image_encoder_{model_name_map[model_name]}.pth"))
+        torch.save(text_encoder.state_dict(), os.path.join(weight_dir, f"clip_text_encoder_{model_name_map[model_name]}.pth"))
+        model_config = {
+            "embed_dim": model.embed_dim,
+            # vision
+            "image_resolution": model.image_resolution,
+            "vision_layers": model.vision_layers,
+            "vision_width": model.vision_width,
+            "vision_patch_size": model.vision_patch_size,
+            # text
+            "context_length": model.context_length,
+            "vocab_size": model.vocab_size,
+            "transformer_width": model.transformer_width,
+            "transformer_heads": model.transformer_heads,
+            "transformer_layers": model.transformer_layers,
+        }
+        image_encoder_config = {
+            "embed_dim": model.embed_dim,
+            "image_resolution": model.image_resolution,
+            "vision_layers": model.vision_layers,
+            "vision_width": model.vision_width,
+            "vision_patch_size": model.vision_patch_size,
+            "vision_heads": model.vision_heads,
+        }
+        text_encoder_config = {
+            "embed_dim": model.embed_dim,
+            "context_length": model.context_length,
+            "vocab_size": model.vocab_size,
+            "transformer_width": model.transformer_width,
+            "transformer_heads": model.transformer_heads,
+            "transformer_layers": model.transformer_layers,
+        }
+        with open(os.path.join(config_dir, f"clip_{model_name_map[model_name]}.json"), "w") as f:
+            json.dump(model_config, f, indent=4)
+        with open(os.path.join(config_dir, f"clip_image_encoder_{model_name_map[model_name]}.json"), "w") as f:
+            json.dump(image_encoder_config, f, indent=4)
+        with open(os.path.join(config_dir, f"clip_text_encoder_{model_name_map[model_name]}.json"), "w") as f:
+            json.dump(text_encoder_config, f, indent=4)
+    print("Done!")

models/clip/_clip/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "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 significant 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+'</w>' 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] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        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('</w>', ' ')
+        return text

models/clip/_clip/text_encoder.py

@@ -0,0 +1,53 @@
+import torch
+from torch import nn, Tensor
+
+from .blocks import LayerNorm, Transformer
+
+
+class CLIPTextEncoder(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        context_length: int,
+        vocab_size: int,
+        transformer_width: int,
+        transformer_heads: int,
+        transformer_layers: int,
+    ) -> None:
+        super().__init__()
+        self.context_length = context_length
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        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))
+
+    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.transformer.resblocks[0].attn.in_proj_weight.dtype
+
+    def forward(self, text: Tensor):
+        x = self.token_embedding(text).type(self.dtype)
+        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 = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return x

models/clip/_clip/utils.py

@@ -0,0 +1,249 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Union, List
+from pkg_resources import packaging
+
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+import torch
+
+from typing import List, Union
+from tqdm import tqdm
+
+from .model import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+            return download_target
+        else:
+            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+        raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if name in _MODELS:
+        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    with open(model_path, 'rb') as opened_file:
+        try:
+            # loading JIT archive
+            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+            state_dict = None
+        except RuntimeError:
+            # loading saved state dict
+            if jit:
+                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+                jit = False
+            state_dict = torch.load(opened_file, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def _node_get(node: torch._C.Node, key: str):
+        """Gets attributes of a node which is polymorphic over return type.
+        
+        From https://github.com/pytorch/pytorch/pull/82628
+        """
+        sel = node.kindOf(key)
+        return getattr(node, sel)(key)
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(_node_get(node, "value")).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if _node_get(inputs[i].node(), "value") == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, 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].
+    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+    """
+    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]
+    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+    else:
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+    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

models/clip/model.py

@@ -0,0 +1,331 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+import numpy as np
+import os
+import math
+from typing import List, Tuple, Union, Optional
+
+from . import _clip
+from ..utils import _init_weights, make_resnet_layers, Bottleneck, BasicBlock
+from .utils import format_count
+
+curr_dir = os.path.abspath(os.path.dirname(__file__))
+
+
+# resnet50: reduction, channels, embed_dim = 32, 2048, 1024
+# resnet101: reduction, channels, embed_dim = 32, 2048, 512
+# resnet50x4: reduction, channels, embed_dim = 32, 2560, 640
+# resnet50x16: reduction, channels, embed_dim = 32, 3072, 768
+# resnet50x64: reduction, channels, embed_dim = 32, 4096, 1024
+# vit_b_32: reduction, channels, embed_dim = 32, 768, 512
+# vit_b_16: reduction, channels, embed_dim = 16, 768, 512
+# vit_l_14: reduction, channels, embed_dim = 14, 1024, 768
+# vit_l_14_336px: reduction, channels, embed_dim = 14, 1024, 768
+
+resnet_backbones = ["resnet50", "resnet101", "resnet50x4", "resnet50x16", "resnet50x64"]
+vit_backbones = ["vit_b_16", "vit_b_32", "vit_l_14", "vit_l_14_336px"]
+
+
+class CLIP_EBC(nn.Module):
+    def __init__(
+        self,
+        backbone: str,
+        bins: List[Tuple[float, float]],
+        anchor_points: List[float],
+        reduction: Optional[int] = None,
+        freeze_text_encoder: bool = True,
+        prompt_type: str = "number",
+        input_size: Optional[int] = None,
+        num_vpt: Optional[int] = None,
+        deep_vpt: Optional[bool] = None,
+        vpt_drop: Optional[float] = None,
+        decoder_block: Optional[nn.Module] = None,
+        decoder_cfg: Optional[List[Union[str, int]]] = None,
+    ) -> None:
+        super().__init__()
+        assert backbone in resnet_backbones + vit_backbones, f"Backbone should be in {resnet_backbones + vit_backbones}, got {backbone}"
+        self.backbone = backbone
+
+        # Image encoder
+        if backbone in resnet_backbones:
+            self.image_encoder = getattr(_clip, f"{backbone}_img")(features_only=True, out_indices=(-1,), reduction=reduction)
+
+        else:
+            assert input_size is not None, "Expected input_size to be an integer, got None."
+            assert num_vpt is not None, "Expected num_vpt to be an integer, got None."
+            assert deep_vpt is not None, "Expected deep_vpt to be a boolean, got None."
+            assert vpt_drop is not None, "Expected vpt_drop to be a float, got None."
+
+            self.image_encoder = getattr(_clip, f"{backbone}_img")(features_only=True, input_size=input_size)
+            self.image_encoder_depth = len(self.image_encoder.transformer.resblocks)
+
+            # Use VPT. Freeze the image encoder.
+            for param in self.image_encoder.parameters():
+                param.requires_grad = False
+
+            self.num_vpt = num_vpt
+            self.deep_vpt = deep_vpt
+
+            patch_size = self.image_encoder.patch_size[0]
+            val = math.sqrt(6. / float(3 * patch_size + self.image_encoder.channels))
+
+            for idx in range(self.image_encoder_depth if self.deep_vpt else 1):
+                setattr(self, f"vpt_{idx}", nn.Parameter(torch.empty(self.num_vpt, self.image_encoder.channels)))
+                nn.init.uniform_(getattr(self, f"vpt_{idx}"), -val, val)
+                setattr(self, f"vpt_drop_{idx}", nn.Dropout(vpt_drop) if vpt_drop > 0 else nn.Identity())
+
+        self.encoder_reduction = self.image_encoder.reduction
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = self.image_encoder.channels
+        self.clip_embed_dim = self.image_encoder.clip_embed_dim
+
+        if decoder_cfg is not None:
+            assert decoder_block is not None, "Expected decoder_block to be a nn.Module, got None."
+            self.image_decoder = make_resnet_layers(decoder_block, decoder_cfg, in_channels=self.channels, expansion=1, dilation=1)
+            self.image_decoder.apply(_init_weights)
+            self.channels = decoder_cfg[-1]
+        else:
+            self.image_decoder = nn.Identity()
+
+        if self.channels != self.clip_embed_dim:
+            self.projection = nn.Conv2d(in_channels=self.channels, out_channels=self.clip_embed_dim, kernel_size=1)
+            self.projection.apply(_init_weights)
+        else:
+            self.projection = nn.Identity()
+
+        # Text encoder
+        assert prompt_type in ["number", "word"], f"Expected prompt_type to be 'number' or 'word', got {prompt_type}"
+        self.prompt_type = prompt_type
+        self.text_encoder = getattr(_clip, f"{backbone}_txt")()
+        self.freeze_text_encoder = freeze_text_encoder
+        if self.freeze_text_encoder:
+            for param in self.text_encoder.parameters():
+                param.requires_grad = False
+
+        self.bins = bins
+        self.anchor_points = torch.tensor(anchor_points, dtype=torch.float32, requires_grad=False).view(1, -1, 1, 1)
+
+        self._get_text_prompts()
+        self._tokenize_text_prompts()
+
+        if self.freeze_text_encoder:
+            self._extract_text_features()
+        else:
+            self.text_features = None
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07), requires_grad=True)
+
+    def _get_text_prompts(self) -> None:
+        bins = [b[0] if b[0] == b[1] else b for b in self.bins]
+        self.text_prompts = [format_count(b, self.prompt_type) for b in bins]
+        print(f"Initialized model with text prompts: {self.text_prompts}")
+
+    def _tokenize_text_prompts(self) -> None:
+        self.text_prompts = _clip.tokenize(self.text_prompts)
+
+    def _extract_text_features(self) -> None:
+        with torch.no_grad():
+            self.text_features = self.text_encoder(self.text_prompts)
+
+    def _prepare_vpt(self, layer: int, batch_size: int, device: torch.device) -> Tensor:
+        if not self.deep_vpt:
+            assert layer == 0, f"Expected layer to be 0 when using Shallow Visual Prompt Tuning, got {layer}"
+
+        vpt = getattr(self, f"vpt_{layer}").to(device)
+        vpt = vpt.unsqueeze(0).expand(batch_size, -1, -1)
+        vpt = getattr(self, f"vpt_drop_{layer}")(vpt)
+        vpt = vpt.permute(1, 0, 2)  # (num_vpt, batch_size, hidden_dim)
+        assert vpt.shape[1] == batch_size, f"Expected the VPT to have the shape [L_vis B C], got {vpt.shape}."
+        return vpt
+
+    def _forward_vpt(self, x: Tensor) -> Tuple[Tensor]:
+        device = x.device
+        batch_size, _, height, width = x.shape
+        num_h_patches, num_w_patches = height // self.image_encoder.patch_size[0], width // self.image_encoder.patch_size[1]
+
+        image_features = self.image_encoder.conv1(x)
+        image_features = image_features.reshape(batch_size, image_features.shape[1], -1)
+        image_features = image_features.permute(0, 2, 1)  # (B, num_patches, C)
+        image_features = torch.cat([
+            self.image_encoder.class_embedding + torch.zeros(batch_size, 1, image_features.shape[-1], dtype=image_features.dtype, device=device),
+            image_features,
+        ], dim=1)  # (B, num_patches + 1, C)
+
+        pos_embedding = self.image_encoder._interpolate_pos_embed(num_h_patches, num_w_patches)
+        image_features = image_features + pos_embedding
+        image_features = self.image_encoder.ln_pre(image_features)
+        image_features = image_features.permute(1, 0, 2)  # (num_patches + 1, B, C)
+        assert image_features.shape[0] == num_h_patches * num_w_patches + 1 and image_features.shape[1] == batch_size, f"Expected image_features to have shape [num_patches + 1, B, C], got {image_features.shape}."
+
+        vpt = self._prepare_vpt(0, batch_size, device)
+        for idx in range(self.image_encoder_depth):
+            # assemble
+            image_features = torch.cat([
+                image_features[:1, :, :],  # CLS token
+                vpt,
+                image_features[1:, :, :],
+            ], dim=0)
+
+            # transformer
+            image_features = self.image_encoder.transformer.resblocks[idx](image_features)
+
+            # disassemble
+            if idx < self.image_encoder_depth - 1:
+                if self.deep_vpt:
+                    vpt = self._prepare_vpt(idx + 1, batch_size, device)
+                else:
+                    vpt = image_features[1: (self.num_vpt + 1), :, :]
+
+            image_features = torch.cat([
+                image_features[:1, :, :],  # CLS token
+                image_features[(self.num_vpt + 1):, :, :],
+            ], dim=0)
+            
+        image_features = image_features.permute(1, 0, 2)  # (B, num_patches + 1, C)
+        image_features = self.image_encoder.ln_post(image_features)
+        image_features = image_features[:, 1:, :].permute(0, 2, 1)  # (B, C, num_patches)
+        image_features = image_features.reshape(batch_size, -1, num_h_patches, num_w_patches)
+        return image_features
+
+    def _forward(self, x: Tensor) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        device = x.device
+
+        x = self.image_encoder(x) if self.backbone in resnet_backbones else self._forward_vpt(x)
+        if self.reduction != self.encoder_reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.image_decoder(x)
+        x = self.projection(x)
+
+        image_features = x.permute(0, 2, 3, 1)  # shape (B, H, W, C)
+        text_features = self.text_encoder(self.text_prompts.to(device)) if self.text_features is None else self.text_features.to(device)  # shape (N, C)
+
+        image_features = F.normalize(image_features, p=2, dim=-1)
+        text_features = F.normalize(text_features, p=2, dim=-1)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits = logit_scale * image_features @ text_features.t()  # (B, H, W, N), logits per image
+        logits = logits.permute(0, 3, 1, 2)  # (B, N, H, W)
+
+        probs = logits.softmax(dim=1)
+        exp = (probs * self.anchor_points.to(x.device)).sum(dim=1, keepdim=True)  # (B, 1, H, W)
+
+        if self.training:
+            return logits, exp
+        else:
+            return exp
+    
+    def forward(self, x: Tensor) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        assert len(x.shape) == 4, f"Expected input to have shape (B C H W), got {x.shape}."
+        if "vit" in self.backbone:
+            image_height, image_width = x.shape[2], x.shape[3]
+            window_height, window_width = self.image_encoder.input_resolution
+
+            if self.training:
+                assert (image_height, image_width) == (window_height, window_width), f"Expected input to have shape ({window_height} {window_width}), got ({image_height} {image_width})."
+                return self._forward(x)
+            
+            elif (image_height, image_width) == (window_height, window_width):  # evaluation, input size = training size
+                return self._forward(x)
+            
+            else:  # evaluation, input_size != training size, use sliding window prediction
+                stride_height, stride_width = window_height, window_width
+                reduction = self.reduction
+
+                num_rows = int(np.ceil((image_height - window_height) / stride_height) + 1)
+                num_cols = int(np.ceil((image_width - window_width) / stride_width) + 1)
+
+                windows = []
+                for i in range(num_rows):
+                    for j in range(num_cols):
+                        x_start, y_start = i * stride_height, j * stride_width
+                        x_end, y_end = x_start + window_height, y_start + window_width
+                        if x_end > image_height:
+                            x_start, x_end = image_height - window_height, image_height
+                        if y_end > image_width:
+                            y_start, y_end = image_width - window_width, image_width
+
+                        window = x[:, :, x_start:x_end, y_start:y_end]
+                        windows.append(window)
+
+                windows = torch.cat(windows, dim=0).to(x.device)  # batched windows, shape: (num_windows, c, h, w)
+
+                preds = self._forward(windows)
+                preds = preds.cpu().detach().numpy()
+
+                # assemble the density map
+                pred_map = np.zeros((preds.shape[1], image_height // reduction, image_width // reduction), dtype=np.float32)
+                count_map = np.zeros((preds.shape[1], image_height // reduction, image_width // reduction), dtype=np.float32)
+                idx = 0
+                for i in range(num_rows):
+                    for j in range(num_cols):
+                        x_start, y_start = i * stride_height, j * stride_width
+                        x_end, y_end = x_start + window_height, y_start + window_width
+                        if x_end > image_height:
+                            x_start, x_end = image_height - window_height, image_height
+                        if y_end > image_width:
+                            y_start, y_end = image_width - window_width, image_width
+
+                        pred_map[:, (x_start // reduction): (x_end // reduction), (y_start // reduction): (y_end // reduction)] += preds[idx, :, :, :]
+                        count_map[:, (x_start // reduction): (x_end // reduction), (y_start // reduction): (y_end // reduction)] += 1.
+                        idx += 1
+
+                pred_map /= count_map  # average the overlapping regions
+                return torch.tensor(pred_map).unsqueeze(0)  # shape: (1, 1, h // reduction, w // reduction)
+        
+        else:
+            return self._forward(x)
+
+
+def _clip_ebc(
+    backbone: str,
+    bins: List[Tuple[float, float]],
+    anchor_points: List[float],
+    reduction: Optional[int] = None,
+    freeze_text_encoder: bool = True,
+    prompt_type: str = "number",
+    input_size: Optional[int] = None,
+    num_vpt: Optional[int] = None,
+    deep_vpt: Optional[bool] = None,
+    vpt_drop: Optional[float] = None,
+    decoder_block: Optional[nn.Module] = None,
+    decoder_cfg: Optional[List[Union[str, int]]] = None
+) -> CLIP_EBC:
+    if backbone in resnet_backbones:
+        decoder_block = Bottleneck
+        if decoder_cfg is None:
+            if backbone == "resnet50":
+                decoder_cfg = [2048]
+            elif backbone == "resnet50x4":
+                decoder_cfg = [1280]
+            elif backbone == "resnet50x16":
+                decoder_cfg = [1536]
+            elif backbone == "resnet50x64":
+                decoder_cfg = [2048]
+            else:  # backbone == "resnet101"
+                decoder_cfg = [2048, 1024]
+    else:
+        decoder_block = BasicBlock
+        if decoder_cfg is None:
+            if backbone == "vit_b_16":
+                decoder_cfg = [768]
+            elif backbone == "vit_b_32":
+                decoder_cfg = [768]
+            else:  # backbone == "vit_l_14"
+                decoder_cfg = [1024]
+
+    return CLIP_EBC(
+        backbone=backbone,
+        bins=bins,
+        anchor_points=anchor_points,
+        reduction=reduction,
+        freeze_text_encoder=freeze_text_encoder,
+        prompt_type=prompt_type,
+        input_size=input_size,
+        num_vpt=num_vpt,
+        deep_vpt=deep_vpt,
+        vpt_drop=vpt_drop,
+        decoder_block=decoder_block,
+        decoder_cfg=decoder_cfg,
+    )

models/clip/utils.py

@@ -0,0 +1,40 @@
+from typing import Union, Tuple
+
+
+num_to_word = {
+    "0": "zero", "1": "one", "2": "two", "3": "three", "4": "four", "5": "five", "6": "six", "7": "seven", "8": "eight", "9": "nine", 
+    "10": "ten", "11": "eleven", "12": "twelve", "13": "thirteen", "14": "fourteen", "15": "fifteen", "16": "sixteen", "17": "seventeen", "18": "eighteen", "19": "nineteen", 
+    "20": "twenty", "21": "twenty-one", "22": "twenty-two", "23": "twenty-three", "24": "twenty-four", "25": "twenty-five", "26": "twenty-six", "27": "twenty-seven", "28": "twenty-eight", "29": "twenty-nine",
+    "30": "thirty", "31": "thirty-one", "32": "thirty-two", "33": "thirty-three", "34": "thirty-four", "35": "thirty-five", "36": "thirty-six", "37": "thirty-seven", "38": "thirty-eight", "39": "thirty-nine",
+    "40": "forty", "41": "forty-one", "42": "forty-two", "43": "forty-three", "44": "forty-four", "45": "forty-five", "46": "forty-six", "47": "forty-seven", "48": "forty-eight", "49": "forty-nine",
+    "50": "fifty", "51": "fifty-one", "52": "fifty-two", "53": "fifty-three", "54": "fifty-four", "55": "fifty-five", "56": "fifty-six", "57": "fifty-seven", "58": "fifty-eight", "59": "fifty-nine",
+    "60": "sixty", "61": "sixty-one", "62": "sixty-two", "63": "sixty-three", "64": "sixty-four", "65": "sixty-five", "66": "sixty-six", "67": "sixty-seven", "68": "sixty-eight", "69": "sixty-nine",
+    "70": "seventy", "71": "seventy-one", "72": "seventy-two", "73": "seventy-three", "74": "seventy-four", "75": "seventy-five", "76": "seventy-six", "77": "seventy-seven", "78": "seventy-eight", "79": "seventy-nine",
+    "80": "eighty", "81": "eighty-one", "82": "eighty-two", "83": "eighty-three", "84": "eighty-four", "85": "eighty-five", "86": "eighty-six", "87": "eighty-seven", "88": "eighty-eight", "89": "eighty-nine",
+    "90": "ninety", "91": "ninety-one", "92": "ninety-two", "93": "ninety-three", "94": "ninety-four", "95": "ninety-five", "96": "ninety-six", "97": "ninety-seven", "98": "ninety-eight", "99": "ninety-nine",
+    "100": "one hundred", "200": "two hundred", "300": "three hundred", "400": "four hundred", "500": "five hundred", "600": "six hundred", "700": "seven hundred", "800": "eight hundred", "900": "nine hundred",
+    "1000": "one thousand"
+}
+
+
+def num2word(num: Union[int, str]) -> str:
+    """
+    Convert the input number to the corresponding English word. For example, 1 -> "one", 2 -> "two", etc.
+    """
+    num = str(int(num))
+    return num_to_word.get(num, num)
+
+
+def format_count(count: Union[float, Tuple[float, float]], prompt_type: str = "word") -> str:
+    if count == 0:
+        return "There is no person." if prompt_type == "word" else "There is 0 person."
+    elif count == 1:
+        return "There is one person." if prompt_type == "word" else "There is 1 person."
+    elif isinstance(count, (int, float)):
+        return f"There are {num2word(int(count))} people." if prompt_type == "word" else f"There are {int(count)} people."
+    elif count[1] == float("inf"):
+        return f"There are more than {num2word(int(count[0]))} people." if prompt_type == "word" else f"There are more than {int(count[0])} people."
+    else:  # count is a tuple of finite numbers
+        left, right = int(count[0]), int(count[1])
+        left, right = num2word(left), num2word(right) if prompt_type == "word" else left, right
+        return f"There are between {left} and {right} people."

models/encoder/__init__.py

@@ -0,0 +1,10 @@
+from .vgg import vgg11, vgg11_bn, vgg13, vgg13_bn, vgg16, vgg16_bn, vgg19, vgg19_bn
+from .vit import vit_b_16, vit_b_32, vit_l_16, vit_l_32, vit_h_14
+from .timm_models import _timm_encoder
+
+
+__all__ = [
+    "vgg11", "vgg11_bn", "vgg13", "vgg13_bn", "vgg16", "vgg16_bn", "vgg19", "vgg19_bn",
+    "vit_b_16", "vit_b_32", "vit_l_16", "vit_l_32", "vit_h_14",
+    "_timm_encoder",
+]

models/encoder/timm_models.py

@@ -0,0 +1,54 @@
+from timm import create_model, list_models
+from torch import nn, Tensor
+import torch.nn.functional as F
+from typing import Optional
+
+from warnings import warn
+
+
+class TIMMEncoder(nn.Module):
+    def __init__(
+        self,
+        backbone: str,
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        assert backbone in list_models(), f"Backbone {backbone} not available in timm"
+        encoder = create_model(backbone, pretrained=True, features_only=True, out_indices=[-1])
+        encoder_reduction = encoder.feature_info.reduction()[-1]
+
+        if reduction <= 16:
+            if "resnet" in backbone:
+                if "resnet18" in backbone or "resnet34" in backbone:
+                    encoder.layer4[0].conv1.stride = (1, 1)
+                    encoder.layer4[0].downsample[0].stride = (1, 1)
+                else:
+                    encoder.layer4[0].conv2.stride = (1, 1)
+                    encoder.layer4[0].downsample[0].stride = (1, 1)
+                encoder_reduction = encoder_reduction // 2
+
+            elif "mobilenetv2" in backbone:
+                encoder.blocks[5][0].conv_dw.stride = (1, 1)
+                encoder_reduction = encoder_reduction // 2
+
+            elif "densenet" in backbone:
+                encoder.features_transition3.pool = nn.Identity()
+                encoder_reduction = encoder_reduction // 2
+
+            else:
+                warn(f"Reduction for {backbone} not handled. Using default reduction of {encoder_reduction}")
+
+        self.encoder = encoder
+        self.encoder_reduction = encoder_reduction
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = self.encoder.feature_info.channels()[-1]
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.encoder(x)[-1]
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        return x
+
+
+def _timm_encoder(backbone: str, reduction: Optional[int] = None) -> TIMMEncoder:
+    return TIMMEncoder(backbone, reduction)

models/encoder/vgg.py

@@ -0,0 +1,69 @@
+from torch import nn, Tensor
+import torch.nn.functional as F
+from torch.hub import load_state_dict_from_url
+from typing import Optional
+
+from ..utils import make_vgg_layers, vgg_cfgs, vgg_urls
+
+
+class VGG(nn.Module):
+    def __init__(
+        self,
+        features: nn.Module,
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self.features = features
+        self.encoder_reduction = 16
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = 512
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        return x
+
+
+def _load_weights(model: VGG, url: str) -> VGG:
+    state_dict = load_state_dict_from_url(url)
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    print("Loading pre-trained weights")
+    if len(missing_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+    if len(unexpected_keys) > 0:
+        print(f"Unexpected keys: {unexpected_keys}")
+    return model
+
+
+def vgg11(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["A"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg11"])
+
+def vgg11_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["A"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg11_bn"])
+
+def vgg13(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["B"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg13"])
+
+def vgg13_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["B"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg13_bn"])
+
+def vgg16(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["D"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg16"])
+
+def vgg16_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["D"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg16_bn"])
+
+def vgg19(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["E"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg19"])
+
+def vgg19_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["E"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg19_bn"])

models/encoder/vit.py

@@ -0,0 +1,526 @@
+import math
+from collections import OrderedDict
+from functools import partial
+from typing import Any, Callable, List, NamedTuple, Optional, Tuple
+
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from torch.hub import load_state_dict_from_url
+from einops import rearrange
+
+from ..utils import Conv2dNormActivation, MLP
+from ..utils import _log_api_usage_once
+
+
+weights = {
+    "vit_b_16": "https://download.pytorch.org/models/vit_b_16-c867db91.pth",
+    "vit_b_32": "https://download.pytorch.org/models/vit_b_32-d86f8d99.pth",
+    "vit_l_16": "https://download.pytorch.org/models/vit_l_16-852ce7e3.pth",
+    "vit_l_32": "https://download.pytorch.org/models/vit_l_32-c7638314.pth",
+    "vit_h_14": "https://download.pytorch.org/models/vit_h_14-6kbcf7eb.pth",
+}
+
+
+class ConvStemConfig(NamedTuple):
+    out_channels: int
+    kernel_size: int
+    stride: int
+    norm_layer: Callable[..., nn.Module] = nn.BatchNorm2d
+    activation_layer: Callable[..., nn.Module] = nn.ReLU
+
+
+class MLPBlock(MLP):
+    """Transformer MLP block."""
+
+    _version = 2
+
+    def __init__(self, in_dim: int, mlp_dim: int, dropout: float):
+        super().__init__(in_dim, [mlp_dim, in_dim], activation_layer=nn.GELU, inplace=None, dropout=dropout)
+
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.normal_(m.bias, std=1e-6)
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        version = local_metadata.get("version", None)
+
+        if version is None or version < 2:
+            # Replacing legacy MLPBlock with MLP. See https://github.com/pytorch/vision/pull/6053
+            for i in range(2):
+                for type in ["weight", "bias"]:
+                    old_key = f"{prefix}linear_{i+1}.{type}"
+                    new_key = f"{prefix}{3*i}.{type}"
+                    if old_key in state_dict:
+                        state_dict[new_key] = state_dict.pop(old_key)
+
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
+
+class EncoderBlock(nn.Module):
+    """Transformer encoder block."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float,
+        attention_dropout: float,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+
+        # Attention block
+        self.ln_1 = norm_layer(hidden_dim)
+        self.self_attention = nn.MultiheadAttention(hidden_dim, num_heads, dropout=attention_dropout, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+
+        # MLP block
+        self.ln_2 = norm_layer(hidden_dim)
+        self.mlp = MLPBlock(hidden_dim, mlp_dim, dropout)
+
+    def forward(self, input: Tensor):
+        torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
+        x = self.ln_1(input)
+        x, _ = self.self_attention(x, x, x, need_weights=False)
+        x = self.dropout(x)
+        x = x + input
+
+        y = self.ln_2(x)
+        y = self.mlp(y)
+        return x + y
+
+
+class Encoder(nn.Module):
+    """Transformer Model Encoder for sequence to sequence translation."""
+    def __init__(
+        self,
+        num_h_patches: int,
+        num_w_patches: int,
+        num_layers: int,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float,
+        attention_dropout: float,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+    ):
+        super().__init__()
+        self.num_h_patches = num_h_patches
+        self.num_w_patches = num_w_patches
+
+        # Note that batch_size is on the first dim because
+        # we have batch_first=True in nn.MultiAttention() by default
+        seq_length = num_h_patches * num_w_patches + 1  # +1 for the class token
+        self.pos_embedding = nn.Parameter(torch.empty(1, seq_length, hidden_dim).normal_(std=0.02))  # from BERT
+        self.dropout = nn.Dropout(dropout)
+        layers: OrderedDict[str, nn.Module] = OrderedDict()
+        for i in range(num_layers):
+            layers[f"encoder_layer_{i}"] = EncoderBlock(
+                num_heads,
+                hidden_dim,
+                mlp_dim,
+                dropout,
+                attention_dropout,
+                norm_layer,
+            )
+        self.layers = nn.Sequential(layers)
+        self.ln = norm_layer(hidden_dim)
+
+    def _get_pos_embedding(self, n_h: int, n_w: int) -> Tensor:
+        if n_h == self.num_h_patches and n_w == self.num_w_patches:
+            return self.pos_embedding
+        else:
+            pos_embedding = self.pos_embedding[:, 1:, :]
+            pos_embedding = rearrange(pos_embedding, "1 (h w) d -> 1 d h w", h=self.num_h_patches, w=self.num_w_patches)
+            pos_embedding = F.interpolate(pos_embedding, size=(n_h, n_w), mode="bicubic")
+            pos_embedding = rearrange(pos_embedding, "1 d h w -> 1 (h w) d")
+            return torch.cat([self.pos_embedding[:, :1, :], pos_embedding], dim=1)
+
+    def forward(self, input: Tensor, n_h: int, n_w: int) -> Tensor:
+        torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
+        input = input + self._get_pos_embedding(n_h, n_w)
+        return self.ln(self.layers(self.dropout(input)))
+
+
+class VisionTransformer(nn.Module):
+    """Vision Transformer as a feature extractor."""
+
+    def __init__(
+        self,
+        image_size: int,
+        patch_size: int,
+        num_layers: int,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float = 0.0,
+        attention_dropout: float = 0.0,
+        # num_classes: int = 1000,  # No need for the classification head as we only need the features
+        reduction: Optional[int] = None,
+        representation_size: Optional[int] = None,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+        conv_stem_configs: Optional[List[ConvStemConfig]] = None,
+    ):
+        super().__init__()
+        _log_api_usage_once(self)
+        torch._assert(image_size % patch_size == 0, "Input shape indivisible by patch size!")
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.hidden_dim = hidden_dim
+        self.mlp_dim = mlp_dim
+        self.attention_dropout = attention_dropout
+        self.dropout = dropout
+        # self.num_classes = num_classes
+        self.representation_size = representation_size
+        self.norm_layer = norm_layer
+
+        if conv_stem_configs is not None:
+            # As per https://arxiv.org/abs/2106.14881
+            seq_proj = nn.Sequential()
+            prev_channels = 3
+            for i, conv_stem_layer_config in enumerate(conv_stem_configs):
+                seq_proj.add_module(
+                    f"conv_bn_relu_{i}",
+                    Conv2dNormActivation(
+                        in_channels=prev_channels,
+                        out_channels=conv_stem_layer_config.out_channels,
+                        kernel_size=conv_stem_layer_config.kernel_size,
+                        stride=conv_stem_layer_config.stride,
+                        norm_layer=conv_stem_layer_config.norm_layer,
+                        activation_layer=conv_stem_layer_config.activation_layer,
+                    ),
+                )
+                prev_channels = conv_stem_layer_config.out_channels
+            seq_proj.add_module(
+                "conv_last", nn.Conv2d(in_channels=prev_channels, out_channels=hidden_dim, kernel_size=1)
+            )
+            self.conv_proj: nn.Module = seq_proj
+        else:
+            self.conv_proj = nn.Conv2d(
+                in_channels=3, out_channels=hidden_dim, kernel_size=patch_size, stride=patch_size
+            )
+
+        seq_length = (image_size // patch_size) ** 2
+
+        # Add a class token
+        self.class_token = nn.Parameter(torch.zeros(1, 1, hidden_dim))
+        seq_length += 1
+
+        self.encoder = Encoder(
+            image_size // patch_size,
+            image_size // patch_size,
+            num_layers,
+            num_heads,
+            hidden_dim,
+            mlp_dim,
+            dropout,
+            attention_dropout,
+            norm_layer,
+        )
+        self.seq_length = seq_length
+
+        # heads_layers: OrderedDict[str, nn.Module] = OrderedDict()
+        # if representation_size is None:
+        #     heads_layers["head"] = nn.Linear(hidden_dim, num_classes)
+        # else:
+        #     heads_layers["pre_logits"] = nn.Linear(hidden_dim, representation_size)
+        #     heads_layers["act"] = nn.Tanh()
+        #     heads_layers["head"] = nn.Linear(representation_size, num_classes)
+
+        # self.heads = nn.Sequential(heads_layers)
+
+        if isinstance(self.conv_proj, nn.Conv2d):
+            # Init the patchify stem
+            fan_in = self.conv_proj.in_channels * self.conv_proj.kernel_size[0] * self.conv_proj.kernel_size[1]
+            nn.init.trunc_normal_(self.conv_proj.weight, std=math.sqrt(1 / fan_in))
+            if self.conv_proj.bias is not None:
+                nn.init.zeros_(self.conv_proj.bias)
+        elif self.conv_proj.conv_last is not None and isinstance(self.conv_proj.conv_last, nn.Conv2d):
+            # Init the last 1x1 conv of the conv stem
+            nn.init.normal_(
+                self.conv_proj.conv_last.weight, mean=0.0, std=math.sqrt(2.0 / self.conv_proj.conv_last.out_channels)
+            )
+            if self.conv_proj.conv_last.bias is not None:
+                nn.init.zeros_(self.conv_proj.conv_last.bias)
+
+        # if hasattr(self.heads, "pre_logits") and isinstance(self.heads.pre_logits, nn.Linear):
+        #     fan_in = self.heads.pre_logits.in_features
+        #     nn.init.trunc_normal_(self.heads.pre_logits.weight, std=math.sqrt(1 / fan_in))
+        #     nn.init.zeros_(self.heads.pre_logits.bias)
+
+        # if isinstance(self.heads.head, nn.Linear):
+        #     nn.init.zeros_(self.heads.head.weight)
+        #     nn.init.zeros_(self.heads.head.bias)
+
+        self.encoder_reduction = self.patch_size
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = hidden_dim
+
+    def _process_input(self, x: Tensor) -> Tuple[Tensor, int, int, int]:
+        # (n, c, h, w) -> (n, hidden_dim, n_h, n_w)
+        x = self.conv_proj(x)
+        n, _, n_h, n_w = x.shape
+        # (n, hidden_dim, n_h, n_w) -> (n, hidden_dim, (n_h * n_w))
+        x = x.reshape(n, self.hidden_dim, n_h * n_w)
+
+        # (n, hidden_dim, (n_h * n_w)) -> (n, (n_h * n_w), hidden_dim)
+        # The self attention layer expects inputs in the format (N, S, E)
+        # where S is the source sequence length, N is the batch size, E is the
+        # embedding dimension
+        x = x.permute(0, 2, 1)
+
+        return x, n, n_h, n_w
+
+    def forward(self, x: Tensor) -> Tensor:
+        # Reshape and permute the input tensor
+        x, n, n_h, n_w = self._process_input(x)
+
+        # Expand the class token to the full batch
+        batch_class_token = self.class_token.expand(n, -1, -1)
+        x = torch.cat([batch_class_token, x], dim=1)
+
+        x = self.encoder(x, n_h, n_w)  # Allows input image to be of any size.
+
+        # Classifier "token" as used by standard language architectures
+        # x = x[:, 0]
+
+        # x = self.heads(x)
+
+        x = x[:, 1:, :]
+        x = rearrange(x, "n (h w) d -> n d h w", h=n_h, w=n_w)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        return x  # To be consistent with timm models
+
+
+def _vision_transformer(
+    patch_size: int,
+    num_layers: int,
+    num_heads: int,
+    hidden_dim: int,
+    mlp_dim: int,
+    weights: str,
+    **kwargs: Any,
+) -> VisionTransformer:
+    image_size = kwargs.pop("image_size", 224)
+
+    model = VisionTransformer(
+        image_size=image_size,
+        patch_size=patch_size,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        hidden_dim=hidden_dim,
+        mlp_dim=mlp_dim,
+        **kwargs,
+    )
+
+    if weights is not None:
+        weights = load_state_dict_from_url(weights, progress=kwargs.get("progress", True))
+        missing_keys, unexpected_keys = model.load_state_dict(weights, strict=False)
+        if len(missing_keys) > 0:
+            print(f"Missing keys: {missing_keys}")
+        if len(unexpected_keys) > 0:
+            print(f"Unexpected keys: {unexpected_keys}")
+
+    return model
+
+
+def interpolate_embeddings(
+    image_size: int,
+    patch_size: int,
+    pos_embedding: Tensor,
+    interpolation_mode: str = "bicubic",
+) -> Tensor:
+    """This function helps interpolate positional embeddings during checkpoint loading,
+    especially when you want to apply a pre-trained model on images with different resolution.
+
+    Args:
+        image_size (int): Image size of the new model.
+        patch_size (int): Patch size of the new model.
+        model_state (OrderedDict[str, Tensor]): State dict of the pre-trained model.
+        interpolation_mode (str): The algorithm used for upsampling. Default: bicubic.
+        reset_heads (bool): If true, not copying the state of heads. Default: False.
+
+    Returns:
+        Tensor: The interpolated positional embedding.
+    """
+    # Shape of pos_embedding is (1, seq_length, hidden_dim)
+    n, seq_length, hidden_dim = pos_embedding.shape
+    if n != 1:
+        raise ValueError(f"Unexpected position embedding shape: {pos_embedding.shape}")
+
+    new_seq_length = (image_size // patch_size) ** 2 + 1
+
+    # Need to interpolate the weights for the position embedding.
+    # We do this by reshaping the positions embeddings to a 2d grid, performing
+    # an interpolation in the (h, w) space and then reshaping back to a 1d grid.
+    if new_seq_length != seq_length:
+        # The class token embedding shouldn't be interpolated, so we split it up.
+        seq_length -= 1
+        new_seq_length -= 1
+        pos_embedding_token = pos_embedding[:, :1, :]
+        pos_embedding_img = pos_embedding[:, 1:, :]
+
+        # (1, seq_length, hidden_dim) -> (1, hidden_dim, seq_length)
+        pos_embedding_img = pos_embedding_img.permute(0, 2, 1)
+        seq_length_1d = int(math.sqrt(seq_length))
+        if seq_length_1d * seq_length_1d != seq_length:
+            raise ValueError(
+                f"seq_length is not a perfect square! Instead got seq_length_1d * seq_length_1d = {seq_length_1d * seq_length_1d } and seq_length = {seq_length}"
+            )
+
+        # (1, hidden_dim, seq_length) -> (1, hidden_dim, seq_l_1d, seq_l_1d)
+        pos_embedding_img = pos_embedding_img.reshape(1, hidden_dim, seq_length_1d, seq_length_1d)
+        new_seq_length_1d = image_size // patch_size
+
+        # Perform interpolation.
+        # (1, hidden_dim, seq_l_1d, seq_l_1d) -> (1, hidden_dim, new_seq_l_1d, new_seq_l_1d)
+        new_pos_embedding_img = nn.functional.interpolate(
+            pos_embedding_img,
+            size=new_seq_length_1d,
+            mode=interpolation_mode,
+        )
+
+        # (1, hidden_dim, new_seq_l_1d, new_seq_l_1d) -> (1, hidden_dim, new_seq_length)
+        new_pos_embedding_img = new_pos_embedding_img.reshape(1, hidden_dim, new_seq_length)
+
+        # (1, hidden_dim, new_seq_length) -> (1, new_seq_length, hidden_dim)
+        new_pos_embedding_img = new_pos_embedding_img.permute(0, 2, 1)
+        new_pos_embedding = torch.cat([pos_embedding_token, new_pos_embedding_img], dim=1)
+
+        return new_pos_embedding
+
+    return pos_embedding
+
+
+def vit_b_16(
+    image_size: int = 224,
+    reduction: int = 16,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=16,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        mlp_dim=3072,
+        weights=weights["vit_b_16"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 16, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+
+
+def vit_b_32(
+    image_size: int = 224,
+    reduction: int = 32,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=32,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        mlp_dim=3072,
+        weights=weights["vit_b_32"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 32, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+
+
+def vit_l_16(
+    image_size: int = 224,
+    reduction: int = 16,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=16,
+        num_layers=24,
+        num_heads=16,
+        hidden_dim=1024,
+        mlp_dim=4096,
+        weights=weights["vit_l_16"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 16, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+
+
+def vit_l_32(
+    image_size: int = 224,
+    reduction: int = 32,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=32,
+        num_layers=24,
+        num_heads=16,
+        hidden_dim=1024,
+        mlp_dim=4096,
+        weights=weights["vit_l_32"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 32, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+
+
+def vit_h_14(
+    image_size: int = 224,
+    reduction: int = 14,
+    **kwargs: Any,
+) -> VisionTransformer:
+    vit = _vision_transformer(
+        patch_size=14,
+        num_layers=32,
+        num_heads=16,
+        hidden_dim=1280,
+        mlp_dim=5120,
+        weights=weights["vit_h_14"],
+        reduction=reduction,
+        **kwargs,
+    )
+    if image_size != 224:
+        vit.image_size = image_size
+        new_pos_embedding = interpolate_embeddings(image_size, 14, vit.state_dict()["encoder.pos_embedding"], "bicubic")
+        vit.encoder.pos_embedding = nn.Parameter(new_pos_embedding, requires_grad=True)
+    return vit
+

models/encoder_decoder/__init__.py

@@ -0,0 +1,17 @@
+from .vgg import vgg11 as vgg11_ae, vgg11_bn as vgg11_bn_ae
+from .vgg import vgg13 as vgg13_ae, vgg13_bn as vgg13_bn_ae
+from .vgg import vgg16 as vgg16_ae, vgg16_bn as vgg16_bn_ae
+from .vgg import vgg19 as vgg19_ae, vgg19_bn as vgg19_bn_ae
+from .resnet import resnet18 as resnet18_ae, resnet34 as resnet34_ae
+from .resnet import resnet50 as resnet50_ae, resnet101 as resnet101_ae, resnet152 as resnet152_ae
+
+from .cannet import cannet, cannet_bn
+from .csrnet import csrnet, csrnet_bn
+
+
+__all__ = [
+    "vgg11_ae", "vgg11_bn_ae", "vgg13_ae", "vgg13_bn_ae", "vgg16_ae", "vgg16_bn_ae", "vgg19_ae", "vgg19_bn_ae",
+    "resnet18_ae", "resnet34_ae", "resnet50_ae", "resnet101_ae", "resnet152_ae",
+    "cannet", "cannet_bn",
+    "csrnet", "csrnet_bn",
+]

models/encoder_decoder/cannet.py

@@ -0,0 +1,85 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+
+from typing import List, Optional
+
+from ..utils import _init_weights
+from .csrnet import CSRNet, csrnet, csrnet_bn
+
+EPS = 1e-6
+
+
+class ContextualModule(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int = 512,
+        sizes: List[int] = [1, 2, 3, 6],
+    ) -> None:
+        super().__init__()
+        self.scales = nn.ModuleList([self.__make_scale__(in_channels, size) for size in sizes])
+        self.bottleneck = nn.Conv2d(in_channels * 2, out_channels, kernel_size=1)
+        self.relu = nn.ReLU(inplace=True)
+        self.weight_net = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def __make_weight__(self, feature: Tensor, scale_feature: Tensor) -> Tensor:
+        weight_feature = feature - scale_feature
+        weight_feature = self.weight_net(weight_feature)
+        return F.sigmoid(weight_feature)
+    
+    def __make_scale__(self, channels: int, size: int) -> nn.Module:
+        return nn.Sequential(
+            nn.AdaptiveAvgPool2d(output_size=(size, size)),
+            nn.Conv2d(channels, channels, kernel_size=1, bias=False),
+        )
+
+    def forward(self, feature: Tensor) -> Tensor:
+        h, w = feature.shape[-2:]
+        multi_scales = [F.interpolate(input=scale(feature), size=(h, w), mode="bilinear") for scale in self.scales]
+        weights = [self.__make_weight__(feature, scale_feature) for scale_feature in multi_scales]
+        multi_scales = sum([multi_scales[i] * weights[i] for i in range(len(weights))]) / (sum(weights) + EPS)
+        overall_features = torch.cat([multi_scales, feature], dim=1)
+        overall_features = self.bottleneck(overall_features)
+        overall_features = self.relu(overall_features)
+        return overall_features
+
+
+class CANNet(nn.Module):
+    def __init__(
+        self,
+        csrnet: CSRNet,
+        sizes: List[int] = [1, 2, 3, 6],
+        reduction: Optional[int] = 8,
+    ) -> None:
+        super().__init__()
+        assert isinstance(csrnet, CSRNet), f"csrnet should be an instance of CSRNet, got {type(csrnet)}."
+        assert isinstance(sizes, (tuple, list)), f"sizes should be a list or tuple, got {type(sizes)}."
+        assert len(sizes) > 0, f"Expected at least one size, got {len(sizes)}."
+        assert all([isinstance(size, int) for size in sizes]), f"Expected all size to be int, got {sizes}."
+        self.sizes = sizes
+        self.encoder_reduction = csrnet.encoder_reduction
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+
+        self.features = csrnet.features
+        self.decoder = csrnet.decoder
+        self.decoder.apply(_init_weights)
+        self.context = ContextualModule(512, 512, self.sizes)
+        self.context.apply(_init_weights)
+
+        self.channels = csrnet.channels
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        x = self.context(x)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.decoder(x)
+        return x
+
+
+def cannet(sizes=[1, 2, 3, 6], reduction: int = 8) -> CANNet:
+    return CANNet(csrnet(), sizes=sizes, reduction=reduction)
+
+def cannet_bn(sizes=[1, 2, 3, 6], reduction: int = 8) -> CANNet:
+    return CANNet(csrnet_bn(), sizes=sizes, reduction=reduction)

models/encoder_decoder/csrnet.py

@@ -0,0 +1,54 @@
+from torch import nn, Tensor
+import torch.nn.functional as F
+from typing import Optional
+
+from ..utils import _init_weights, make_vgg_layers, vgg_urls
+from .vgg import _load_weights
+
+EPS = 1e-6
+
+
+encoder_cfg = [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512]
+decoder_cfg = [512, 512, 512, 256, 128, 64]
+
+
+class CSRNet(nn.Module):
+    def __init__(
+        self,
+        features: nn.Module,
+        decoder: nn.Module,
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self.features = features
+        self.features.apply(_init_weights)
+        self.decoder = decoder
+        self.decoder.apply(_init_weights)
+
+        self.encoder_reduction = 8
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = 64
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.decoder(x)
+        return x
+
+
+def csrnet(reduction: int = 8) -> CSRNet:
+    model = CSRNet(
+        make_vgg_layers(encoder_cfg, in_channels=3, batch_norm=False, dilation=1),
+        make_vgg_layers(decoder_cfg, in_channels=512, batch_norm=False, dilation=2),
+        reduction=reduction
+    )
+    return _load_weights(model, vgg_urls["vgg16"])
+
+def csrnet_bn(reduction: int = 8) -> CSRNet:
+    model = CSRNet(
+        make_vgg_layers(encoder_cfg, in_channels=3, batch_norm=True, dilation=1),
+        make_vgg_layers(decoder_cfg, in_channels=512, batch_norm=True, dilation=2),
+        reduction=reduction
+    )
+    return _load_weights(model, vgg_urls["vgg16"])

models/encoder_decoder/resnet.py

@@ -0,0 +1,95 @@
+from torch import nn, Tensor
+import torch.nn.functional as F
+import timm
+from typing import Union, Optional
+
+from ..utils import BasicBlock, Bottleneck, make_resnet_layers
+from ..utils import _init_weights
+
+
+model_configs = {
+    "resnet18.tv_in1k": {
+        "decoder_channels": [512, 256, 128],
+    },
+    "resnet34.tv_in1k": {
+        "decoder_channels": [512, 256, 128],
+    },
+    "resnet50.tv_in1k": {
+        "decoder_channels": [512, 256, 256, 128],
+    },
+    "resnet101.tv_in1k": {
+        "decoder_channels": [512, 512, 256, 256, 128],
+    },
+    "resnet152.tv_in1k": {
+        "decoder_channels": [512, 512, 512, 256, 256, 128],
+    },
+}
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        decoder_block: Union[BasicBlock, Bottleneck],
+        backbone: str = "resnet34.tv_in1k",
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        assert backbone in model_configs.keys(), f"Backbone should be in {model_configs.keys()}"
+        config = model_configs[backbone]
+        encoder = timm.create_model(backbone, pretrained=True, features_only=True, out_indices=(-1,))
+        encoder_reduction = encoder.feature_info.reduction()[-1]
+
+        if reduction <= 16:
+            if "resnet18" in backbone or "resnet34" in backbone:
+                encoder.layer4[0].conv1.stride = (1, 1)
+                encoder.layer4[0].downsample[0].stride = (1, 1)
+            else:
+                encoder.layer4[0].conv2.stride = (1, 1)
+                encoder.layer4[0].downsample[0].stride = (1, 1)
+            encoder_reduction = encoder_reduction // 2
+
+        self.encoder = encoder
+        self.encoder_reduction = encoder_reduction
+
+        encoder_out_channels = self.encoder.feature_info.channels()[-1]
+
+        decoder_channels = config["decoder_channels"]
+        self.decoder = make_resnet_layers(
+            block=decoder_block,
+            cfg=decoder_channels,
+            in_channels=encoder_out_channels,
+            dilation=1,
+            expansion=1,
+        )
+        self.decoder.apply(_init_weights)
+
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = decoder_channels[-1]
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.encoder(x)[-1]
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.decoder(x)
+
+        return x
+
+
+def resnet18(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=BasicBlock, backbone="resnet18.tv_in1k", reduction=reduction)
+
+
+def resnet34(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=BasicBlock, backbone="resnet34.tv_in1k", reduction=reduction)
+
+
+def resnet50(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=Bottleneck, backbone="resnet50.tv_in1k", reduction=reduction)
+
+
+def resnet101(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=Bottleneck, backbone="resnet101.tv_in1k", reduction=reduction)
+
+
+def resnet152(reduction: int = 32) -> ResNet:
+    return ResNet(decoder_block=Bottleneck, backbone="resnet152.tv_in1k", reduction=reduction)

models/encoder_decoder/vgg.py

@@ -0,0 +1,85 @@
+# The model used in the paper Distribution Matching for Crowd Counting.
+# Code adapted from https://github.com/cvlab-stonybrook/DM-Count/blob/master/models.py
+from torch import nn, Tensor
+import torch.nn.functional as F
+from torch.hub import load_state_dict_from_url
+from typing import Optional
+
+from ..utils import make_vgg_layers, vgg_cfgs, vgg_urls
+from ..utils import _init_weights
+
+
+
+class VGG(nn.Module):
+    def __init__(
+        self,
+        features: nn.Module,
+        reduction: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self.features = features
+        self.reg_layer = nn.Sequential(
+            nn.Conv2d(512, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+        )
+
+        self.reg_layer.apply(_init_weights)
+        # Remove the density layer, as the output from this model is not final and will be further processed.
+        # self.density_layer = nn.Sequential(nn.Conv2d(128, 1, 1), nn.ReLU())
+        self.encoder_reduction = 16
+        self.reduction = self.encoder_reduction if reduction is None else reduction
+        self.channels = 128
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        if self.encoder_reduction != self.reduction:
+            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
+        x = self.reg_layer(x)
+        # x = self.density_layer(x)
+        return x
+
+
+def _load_weights(model: VGG, url: str) -> VGG:
+    state_dict = load_state_dict_from_url(url)
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    print("Loading pre-trained weights")
+    if len(missing_keys) > 0:
+        print(f"Missing keys: {missing_keys}")
+    if len(unexpected_keys) > 0:
+        print(f"Unexpected keys: {unexpected_keys}")
+    return model
+
+
+def vgg11(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["A"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg11"])
+
+def vgg11_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["A"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg11_bn"])
+
+def vgg13(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["B"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg13"])
+
+def vgg13_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["B"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg13_bn"])
+
+def vgg16(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["D"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg16"])
+
+def vgg16_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["D"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg16_bn"])
+
+def vgg19(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["E"]), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg19"])
+
+def vgg19_bn(reduction: int = 8) -> VGG:
+    model = VGG(make_vgg_layers(vgg_cfgs["E"], batch_norm=True), reduction=reduction)
+    return _load_weights(model, vgg_urls["vgg19_bn"])

models/model.py

@@ -0,0 +1,112 @@
+import torch
+from torch import nn, Tensor
+import os
+from typing import List, Tuple, Union, Callable
+from functools import partial
+
+from .utils import _init_weights
+
+from . import encoder
+from . import encoder_decoder
+from .encoder import _timm_encoder
+
+
+curr_dir = os.path.abspath(os.path.dirname(__file__))
+
+
+class Regressor(nn.Module):
+    def __init__(self, backbone: nn.Module) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.reduction = backbone.reduction
+
+        self.regressor = nn.Sequential(
+            nn.Conv2d(backbone.channels, 1, kernel_size=1),
+            nn.ReLU(inplace=True),
+        )
+        self.regressor.apply(_init_weights)
+        self.bins = None
+        self.anchor_points = None
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.backbone(x)
+        x = self.regressor(x)
+        return x
+
+
+class Classifier(nn.Module):
+    def __init__(
+        self,
+        backbone: nn.Module,
+        bins: List[Tuple[float, float]],
+        anchor_points: List[float],
+    ) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.reduction = backbone.reduction
+
+        assert len(bins) == len(anchor_points), f"Expected bins and anchor_points to have the same length, got {len(bins)} and {len(anchor_points)}"
+        assert all(len(b) == 2 for b in bins), f"Expected bins to be a list of tuples of length 2, got {bins}"
+        assert all(bin[0] <= p <= bin[1] for bin, p in zip(bins, anchor_points)), f"Expected anchor_points to be within the range of the corresponding bin, got {bins} and {anchor_points}"
+
+        self.bins = bins
+        self.anchor_points = torch.tensor(anchor_points, dtype=torch.float32, requires_grad=False).view(1, -1, 1, 1)
+
+        if backbone.channels > 512:
+            self.classifier = nn.Sequential(
+                nn.Conv2d(backbone.channels, 512, kernel_size=1),  # serves as a linear layer for feature vectors at each pixel
+                nn.ReLU(inplace=True),
+                nn.Conv2d(512, len(self.bins), kernel_size=1),
+            )
+        else:
+            self.classifier = nn.Conv2d(backbone.channels, len(self.bins), kernel_size=1)
+
+        self.classifier.apply(_init_weights)
+
+    def forward(self, x: Tensor) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        x = self.backbone(x)
+        x = self.classifier(x)  # shape (B, C, H, W), where C = len(bins), x is the logits
+
+        probs = x.softmax(dim=1)  # shape (B, C, H, W)
+        exp = (probs * self.anchor_points.to(x.device)).sum(dim=1, keepdim=True)  # shape (B, 1, H, W)
+        if self.training:
+            return x, exp
+        else:
+            return exp
+
+
+def _get_backbone(backbone: str, input_size: int, reduction: int) -> Callable:
+    assert "clip" not in backbone, f"This function does not support CLIP model, got {backbone}"
+
+    if backbone in ["vit_b_16", "vit_b_32", "vit_l_16", "vit_l_32", "vit_h_14"]:
+        return partial(getattr(encoder, backbone), image_size=input_size, reduction=reduction)
+    elif backbone in ["vgg11", "vgg11_bn", "vgg13", "vgg13_bn", "vgg16", "vgg16_bn", "vgg19", "vgg19_bn"]:
+        return partial(getattr(encoder, backbone), reduction=reduction)
+    elif backbone in ["vgg11_ae", "vgg11_bn_ae", "vgg13_ae", "vgg13_bn_ae", "vgg16_ae", "vgg16_bn_ae", "vgg19_ae", "vgg19_bn_ae"]:
+        return partial(getattr(encoder_decoder, backbone), reduction=reduction)
+    elif backbone in ["resnet18_ae", "resnet34_ae", "resnet50_ae", "resnet101_ae", "resnet152_ae"]:
+        return partial(getattr(encoder_decoder, backbone), reduction=reduction)
+    elif backbone in ["cannet", "cannet_bn", "csrnet", "csrnet_bn"]:
+        return partial(getattr(encoder_decoder, backbone), reduction=reduction)
+    else:
+        return partial(_timm_encoder, backbone=backbone, reduction=reduction)
+
+
+def _regressor(
+    backbone: str,
+    input_size: int,
+    reduction: int,
+) -> Regressor:
+    backbone = _get_backbone(backbone.lower(), input_size, reduction)
+    return Regressor(backbone())
+
+
+def _classifier(
+    backbone: nn.Module,
+    input_size: int,
+    reduction: int,
+    bins: List[Tuple[float, float]],
+    anchor_points: List[float],
+) -> Classifier:
+    backbone = _get_backbone(backbone.lower(), input_size, reduction)
+    return Classifier(backbone(), bins, anchor_points)

models/utils.py

@@ -0,0 +1,444 @@
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+from functools import partial
+from typing import Callable, Optional, Sequence, Tuple, Union, Any, List, TypeVar, List
+from types import FunctionType
+from itertools import repeat
+import warnings
+import os
+from collections.abc import Iterable
+
+V = TypeVar("V")
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+
+
+vgg_urls = {
+    "vgg11": "https://download.pytorch.org/models/vgg11-8a719046.pth",
+    "vgg11_bn": "https://download.pytorch.org/models/vgg11_bn-6002323d.pth",
+    "vgg13": "https://download.pytorch.org/models/vgg13-19584684.pth",
+    "vgg13_bn": "https://download.pytorch.org/models/vgg13_bn-abd245e5.pth",
+    "vgg16": "https://download.pytorch.org/models/vgg16-397923af.pth",
+    "vgg16_bn": "https://download.pytorch.org/models/vgg16_bn-6c64b313.pth",
+    "vgg19": "https://download.pytorch.org/models/vgg19-dcbb9e9d.pth",
+    "vgg19_bn": "https://download.pytorch.org/models/vgg19_bn-c79401a0.pth",
+}
+
+vgg_cfgs = {
+    "A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512],
+    "B": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512],
+    "D": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512],
+    "E": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512]
+}
+
+
+def _log_api_usage_once(obj: Any) -> None:
+
+    """
+    Logs API usage(module and name) within an organization.
+    In a large ecosystem, it's often useful to track the PyTorch and
+    TorchVision APIs usage. This API provides the similar functionality to the
+    logging module in the Python stdlib. It can be used for debugging purpose
+    to log which methods are used and by default it is inactive, unless the user
+    manually subscribes a logger via the `SetAPIUsageLogger method <https://github.com/pytorch/pytorch/blob/eb3b9fe719b21fae13c7a7cf3253f970290a573e/c10/util/Logging.cpp#L114>`_.
+    Please note it is triggered only once for the same API call within a process.
+    It does not collect any data from open-source users since it is no-op by default.
+    For more information, please refer to
+    * PyTorch note: https://pytorch.org/docs/stable/notes/large_scale_deployments.html#api-usage-logging;
+    * Logging policy: https://github.com/pytorch/vision/issues/5052;
+
+    Args:
+        obj (class instance or method): an object to extract info from.
+    """
+    module = obj.__module__
+    if not module.startswith("torchvision"):
+        module = f"torchvision.internal.{module}"
+    name = obj.__class__.__name__
+    if isinstance(obj, FunctionType):
+        name = obj.__name__
+    torch._C._log_api_usage_once(f"{module}.{name}")
+
+
+def _make_ntuple(x: Any, n: int) -> Tuple[Any, ...]:
+    """
+    Make n-tuple from input x. If x is an iterable, then we just convert it to tuple.
+    Otherwise, we will make a tuple of length n, all with value of x.
+    reference: https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/utils.py#L8
+
+    Args:
+        x (Any): input value
+        n (int): length of the resulting tuple
+    """
+    if isinstance(x, Iterable):
+        return tuple(x)
+    return tuple(repeat(x, n))
+
+
+class ConvNormActivation(torch.nn.Sequential):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, ...]] = 3,
+        stride: Union[int, Tuple[int, ...]] = 1,
+        padding: Optional[Union[int, Tuple[int, ...], str]] = None,
+        groups: int = 1,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        dilation: Union[int, Tuple[int, ...]] = 1,
+        inplace: Optional[bool] = True,
+        bias: Optional[bool] = None,
+        conv_layer: Callable[..., torch.nn.Module] = torch.nn.Conv2d,
+    ) -> None:
+
+        if padding is None:
+            if isinstance(kernel_size, int) and isinstance(dilation, int):
+                padding = (kernel_size - 1) // 2 * dilation
+            else:
+                _conv_dim = len(kernel_size) if isinstance(kernel_size, Sequence) else len(dilation)
+                kernel_size = _make_ntuple(kernel_size, _conv_dim)
+                dilation = _make_ntuple(dilation, _conv_dim)
+                padding = tuple((kernel_size[i] - 1) // 2 * dilation[i] for i in range(_conv_dim))
+        if bias is None:
+            bias = norm_layer is None
+
+        layers = [
+            conv_layer(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                dilation=dilation,
+                groups=groups,
+                bias=bias,
+            )
+        ]
+
+        if norm_layer is not None:
+            layers.append(norm_layer(out_channels))
+
+        if activation_layer is not None:
+            params = {} if inplace is None else {"inplace": inplace}
+            layers.append(activation_layer(**params))
+        super().__init__(*layers)
+        _log_api_usage_once(self)
+        self.out_channels = out_channels
+
+        if self.__class__ == ConvNormActivation:
+            warnings.warn(
+                "Don't use ConvNormActivation directly, please use Conv2dNormActivation and Conv3dNormActivation instead."
+            )
+
+
+class Conv2dNormActivation(ConvNormActivation):
+    """
+    Configurable block used for Convolution2d-Normalization-Activation blocks.
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the Convolution-Normalization-Activation block
+        kernel_size: (int, optional): Size of the convolving kernel. Default: 3
+        stride (int, optional): Stride of the convolution. Default: 1
+        padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in which case it will be calculated as ``padding = (kernel_size - 1) // 2 * dilation``
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolution layer. If ``None`` this layer won't be used. Default: ``torch.nn.BatchNorm2d``
+        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer won't be used. Default: ``torch.nn.ReLU``
+        dilation (int): Spacing between kernel elements. Default: 1
+        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
+        bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``.
+
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int]] = 3,
+        stride: Union[int, Tuple[int, int]] = 1,
+        padding: Optional[Union[int, Tuple[int, int], str]] = None,
+        groups: int = 1,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        dilation: Union[int, Tuple[int, int]] = 1,
+        inplace: Optional[bool] = True,
+        bias: Optional[bool] = None,
+    ) -> None:
+
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            groups,
+            norm_layer,
+            activation_layer,
+            dilation,
+            inplace,
+            bias,
+            torch.nn.Conv2d,
+        )
+
+
+class MLP(torch.nn.Sequential):
+    """This block implements the multi-layer perceptron (MLP) module.
+
+    Args:
+        in_channels (int): Number of channels of the input
+        hidden_channels (List[int]): List of the hidden channel dimensions
+        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the linear layer. If ``None`` this layer won't be used. Default: ``None``
+        activation_layer (Callable[..., torch.nn.Module], optional): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the linear layer. If ``None`` this layer won't be used. Default: ``torch.nn.ReLU``
+        inplace (bool, optional): Parameter for the activation layer, which can optionally do the operation in-place.
+            Default is ``None``, which uses the respective default values of the ``activation_layer`` and Dropout layer.
+        bias (bool): Whether to use bias in the linear layer. Default ``True``
+        dropout (float): The probability for the dropout layer. Default: 0.0
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: List[int],
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = None,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        inplace: Optional[bool] = None,
+        bias: bool = True,
+        dropout: float = 0.0,
+    ):
+        # The addition of `norm_layer` is inspired from the implementation of TorchMultimodal:
+        # https://github.com/facebookresearch/multimodal/blob/5dec8a/torchmultimodal/modules/layers/mlp.py
+        params = {} if inplace is None else {"inplace": inplace}
+
+        layers = []
+        in_dim = in_channels
+        for hidden_dim in hidden_channels[:-1]:
+            layers.append(torch.nn.Linear(in_dim, hidden_dim, bias=bias))
+            if norm_layer is not None:
+                layers.append(norm_layer(hidden_dim))
+            layers.append(activation_layer(**params))
+            layers.append(torch.nn.Dropout(dropout, **params))
+            in_dim = hidden_dim
+
+        layers.append(torch.nn.Linear(in_dim, hidden_channels[-1], bias=bias))
+        layers.append(torch.nn.Dropout(dropout, **params))
+
+        super().__init__(*layers)
+        _log_api_usage_once(self)
+
+
+def conv3x3(
+    in_channels: int,
+    out_channels: int,
+    stride: int = 1,
+    groups: int = 1,
+    dilation: int = 1,
+) -> nn.Conv2d:
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1x1(in_channels: int, out_channels: int, stride: int = 1) -> nn.Conv2d:
+    """1x1 convolution"""
+    return nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        stride: int = 1,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(in_channels, out_channels, stride)
+        self.bn1 = norm_layer(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(out_channels, out_channels)
+        self.bn2 = norm_layer(out_channels)
+        self.stride = stride
+        if in_channels != out_channels:
+            self.downsample = nn.Sequential(
+                conv1x1(in_channels, out_channels),
+                nn.BatchNorm2d(out_channels),
+            )
+        else:
+            self.downsample = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        out += self.downsample(identity)
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition" https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        stride: int = 1,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        expansion: int = 4,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(out_channels * (base_width / 64.0)) * groups
+        self.expansion = expansion
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(in_channels, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, out_channels * self.expansion)
+        self.bn3 = norm_layer(out_channels * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.stride = stride
+        if in_channels != out_channels:
+            self.downsample = nn.Sequential(
+                conv1x1(in_channels, out_channels),
+                nn.BatchNorm2d(out_channels),
+            )
+        else:
+            self.downsample = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        out += self.downsample(identity)
+        out = self.relu(out)
+
+        return out
+    
+
+def _init_weights(model: nn.Module) -> None:
+    for m in model.modules():
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0.)
+        elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+            nn.init.constant_(m.weight, 1.)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0.)
+        elif isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=0.01)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0.)
+
+
+class Upsample(nn.Module):
+    def __init__(
+        self,
+        size: Union[int, Tuple[int, int]] = None,
+        scale_factor: Union[float, Tuple[float, float]] = None,
+        mode: str = "nearest",
+        align_corners: bool = False,
+        antialias: bool = False,
+    ) -> None:
+        super().__init__()
+        self.interpolate = partial(
+            F.interpolate,
+            size=size,
+            scale_factor=scale_factor,
+            mode=mode,
+            align_corners=align_corners,
+            antialias=antialias,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.interpolate(x)
+
+
+def make_vgg_layers(cfg: List[Union[str, int]], in_channels: int = 3, batch_norm: bool = False, dilation: int = 1) -> nn.Sequential:
+    layers = []
+    for v in cfg:
+        if v == "M":
+            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
+        elif v == "U":
+            layers += [Upsample(scale_factor=2, mode="bilinear")]
+        else:
+            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=dilation, dilation=dilation)
+            if batch_norm:
+                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
+            else:
+                layers += [conv2d, nn.ReLU(inplace=True)]
+            in_channels = v
+    return nn.Sequential(*layers)
+
+
+def make_resnet_layers(
+    block: Union[BasicBlock, Bottleneck],
+    cfg: List[Union[int, str]],
+    in_channels: int,
+    dilation: int = 1,
+    expansion: int = 1,
+) -> nn.Sequential:
+    layers = []
+    for v in cfg:
+        if v == "U":
+            layers.append(Upsample(scale_factor=2, mode="bilinear"))
+        else:
+            layers.append(block(
+                in_channels=in_channels,
+                out_channels=v,
+                dilation=dilation,
+                expansion=expansion,
+            ))
+            in_channels = v
+
+    layers = nn.Sequential(*layers)
+    layers.apply(_init_weights)
+    return layers

preprocess.py

@@ -0,0 +1,458 @@
+import os
+from glob import glob
+from scipy.io import loadmat
+import cv2
+from argparse import ArgumentParser
+from tqdm import tqdm
+import numpy as np
+from typing import Tuple, Union, Optional
+from warnings import warn
+
+from datasets import standardize_dataset_name
+
+
+def _calc_size(
+    img_w: int,
+    img_h: int,
+    min_size: int,
+    max_size: int,
+    base: int = 32
+) -> Union[Tuple[int, int], None]:
+    """
+    This function generates a new size for an image while keeping the aspect ratio. The new size should be within the given range (min_size, max_size).
+
+    Args:
+        img_w (int): The width of the image.
+        img_h (int): The height of the image.
+        min_size (int): The minimum size of the edges of the image.
+        max_size (int): The maximum size of the edges of the image.
+    """
+    assert min_size % base == 0, f"min_size ({min_size}) must be a multiple of {base}"
+    if max_size != float("inf"):
+        assert max_size % base == 0, f"max_size ({max_size}) must be a multiple of {base} if provided"
+
+    assert min_size <= max_size, f"min_size ({min_size}) must be less than or equal to max_size ({max_size})"
+
+    aspect_ratios = (img_w / img_h, img_h / img_w)
+    if min_size / max_size <= min(aspect_ratios) <= max(aspect_ratios) <= max_size / min_size:  # possible to resize and preserve the aspect ratio
+        if min_size <= min(img_w, img_h) <= max(img_w, img_h) <= max_size:  # already within the range, no need to resize
+            ratio = 1.
+        elif min(img_w, img_h) < min_size:  # smaller than the minimum size, resize to the minimum size
+            ratio = min_size / min(img_w, img_h)
+        else:  # larger than the maximum size, resize to the maximum size
+            ratio = max_size / max(img_w, img_h)
+
+        new_w, new_h = int(round(img_w * ratio / base) * base), int(round(img_h * ratio / base) * base)
+        new_w = max(min_size, min(max_size, new_w))
+        new_h = max(min_size, min(max_size, new_h))
+        return new_w, new_h
+
+    else:  # impossible to resize and preserve the aspect ratio
+        msg = f"Impossible to resize {img_w}x{img_h} image while preserving the aspect ratio to a size within the range ({min_size}, {max_size}). Will not limit the maximum size."
+        warn(msg)
+        return _calc_size(img_w, img_h, min_size, float("inf"), base)
+
+
+def _generate_random_indices(
+    total_size: int,
+    out_dir: str,
+) -> None:
+    """
+    Generate randomly selected indices for labelled data in semi-supervised learning.
+    """
+    rng = np.random.default_rng(42)
+    for percent in [0.01, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]:
+        num_select = int(total_size * percent)
+        selected = rng.choice(total_size, num_select, replace=False)
+        selected.sort()
+        selected = selected.tolist()
+        with open(os.path.join(out_dir, f"{int(percent * 100)}%.txt"), "w") as f:
+            for i in selected:
+                f.write(f"{i}\n")
+
+
+def _resize(image: np.ndarray, label: np.ndarray, min_size: int, max_size: int) -> Tuple[np.ndarray, np.ndarray, bool]:
+    image_h, image_w, _ = image.shape
+    new_size = _calc_size(image_w, image_h, min_size, max_size)
+    if new_size is None:
+        return image, label, False
+    else:
+        new_w, new_h = new_size
+        image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_CUBIC) if (new_w, new_h) != (image_w, image_h) else image
+        label = label * np.array([[new_w / image_w, new_h / image_h]]) if len(label) > 0 and (new_w, new_h) != (image_w, image_h) else label
+        return image, label, True
+
+
+def _preprocess(
+    dataset: str,
+    data_src_dir: str,
+    data_dst_dir: str,
+    min_size: int,
+    max_size: int,
+    generate_npy: bool = False
+) -> None:
+    """
+    This function organizes the data into the following structure:
+    data_dst_dir
+    ├── train
+    │   ├── images
+    │   │   ├── 0001.jpg
+    │   │   ├── 0002.jpg
+    │   │   ├── ...
+    │   │   images_npy
+    │   │   ├── 0001.npy
+    │   │   ├── 0002.npy
+    │   │   ├── ...
+    │   ├── labels
+    │   │   ├── 0001.npy
+    │   │   ├── 0002.npy
+    │   │   ├── ...
+    │   ├── 0.01%.txt
+    │   ├── 0.05%.txt
+    │   ├── ...
+    ├── val
+    │   ├── images
+    │   │   ├── 0001.jpg
+    │   │   ├── 0002.jpg
+    │   │   ├── ...
+    │   │   images_npy
+    │   │   ├── 0001.npy
+    │   │   ├── 0002.npy
+    │   │   ├── ...
+    │   ├── labels
+    │   │   ├── 0001.npy
+    │   │   ├── 0002.npy
+    │   │   ├── ...
+    """
+    dataset = standardize_dataset_name(dataset)
+    assert os.path.isdir(data_src_dir), f"{data_src_dir} does not exist"
+    os.makedirs(data_dst_dir, exist_ok=True)
+    print(f"Pre-processing {dataset} dataset...")
+    if dataset in ["sha", "shb"]:
+        _shanghaitech(data_src_dir, data_dst_dir, min_size, max_size, generate_npy)
+
+    elif dataset == "nwpu":
+        _nwpu(data_src_dir, data_dst_dir, min_size, max_size, generate_npy)
+
+    elif dataset == "qnrf":
+        _qnrf(data_src_dir, data_dst_dir, min_size, max_size, generate_npy)
+    
+    else:  # dataset == "jhu"
+        _jhu(data_src_dir, data_dst_dir, min_size, max_size, generate_npy)
+
+
+def _resize_and_save(
+    image: np.ndarray,
+    name: str,
+    image_dst_dir: str,
+    generate_npy: bool,
+    label: Optional[np.ndarray] = None,
+    label_dst_dir: Optional[str] = None,
+    min_size: Optional[int] = None,
+    max_size: Optional[int] = None,
+) -> None:
+    os.makedirs(image_dst_dir, exist_ok=True)
+
+    if label is not None:
+        assert label_dst_dir is not None, "label_dst_dir must be provided if label is provided"
+        os.makedirs(label_dst_dir, exist_ok=True)
+
+    image_dst_path = os.path.join(image_dst_dir, f"{name}.jpg")
+
+    if label is not None:
+        label_dst_path = os.path.join(label_dst_dir, f"{name}.npy")
+    else:
+        label = np.array([])
+        label_dst_path = None
+
+    if min_size is not None:
+        assert max_size is not None, f"max_size must be provided if min_size is provided, got {max_size}"
+        image, label, success = _resize(image, label, min_size, max_size)
+        if not success:
+            print(f"image: {image_dst_path} is not resized")
+
+    cv2.imwrite(image_dst_path, image)
+
+    if label_dst_path is not None:
+        np.save(label_dst_path, label)
+
+    if generate_npy:
+        image_npy_dst_path = os.path.join(image_dst_dir, f"{name}.npy")
+        image_npy = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # convert to RGB
+        image_npy = np.transpose(image_npy, (2, 0, 1))  # HWC to CHW
+        # Don't normalize the image. Keep it as np.uint8 to save space.
+        # image_npy = image_npy.astype(np.float32) / 255.  # normalize to [0, 1]
+        np.save(image_npy_dst_path, image_npy)
+
+
+def _shanghaitech(
+    data_src_dir: str,
+    data_dst_dir: str,
+    min_size: int,
+    max_size: int,
+    generate_npy: bool = False
+) -> None:
+    for split in ["train", "val"]:
+        generate_npy = generate_npy and split == "train"
+        print(f"Processing {split}...")
+        if split == "train":
+            image_src_dir = os.path.join(data_src_dir, "train_data", "images")
+            label_src_dir = os.path.join(data_src_dir, "train_data", "ground-truth")
+            image_src_paths = glob(os.path.join(image_src_dir, "*.jpg"))
+            label_src_paths = glob(os.path.join(label_src_dir, "*.mat"))
+            assert len(image_src_paths) == len(label_src_paths) in [300, 400], f"Expected 300 (part_A) or 400 (part_B) images and labels, got {len(image_src_paths)} images and {len(label_src_paths)} labels"
+        else:
+            image_src_dir = os.path.join(data_src_dir, "test_data", "images")
+            label_src_dir = os.path.join(data_src_dir, "test_data", "ground-truth")
+            image_src_paths = glob(os.path.join(image_src_dir, "*.jpg"))
+            label_src_paths = glob(os.path.join(label_src_dir, "*.mat"))
+            assert len(image_src_paths) == len(label_src_paths) in [182, 316], f"Expected 182 (part_A) or 316 (part_B) images and labels, got {len(image_src_paths)} images and {len(label_src_paths)} labels"
+
+        sort_key = lambda x: int((os.path.basename(x).split(".")[0]).split("_")[-1])
+        image_src_paths.sort(key=sort_key)
+        label_src_paths.sort(key=sort_key)
+
+        image_dst_dir = os.path.join(data_dst_dir, split, "images")
+        label_dst_dir = os.path.join(data_dst_dir, split, "labels")
+        os.makedirs(image_dst_dir, exist_ok=True)
+        os.makedirs(label_dst_dir, exist_ok=True)
+
+        size = len(str(len(image_src_paths)))
+        for i, (image_src_path, label_src_path) in tqdm(enumerate(zip(image_src_paths, label_src_paths)), total=len(image_src_paths)):
+            image_id = int((os.path.basename(image_src_path).split(".")[0]).split("_")[-1])
+            label_id = int((os.path.basename(label_src_path).split(".")[0]).split("_")[-1])
+            assert image_id == label_id, f"Expected image id {image_id} to match label id {label_id}"
+            name = f"{(i + 1):0{size}d}"
+            image = cv2.imread(image_src_path)
+            label = loadmat(label_src_path)["image_info"][0][0][0][0][0]
+            _resize_and_save(
+                image=image,
+                label=label,
+                name=name,
+                image_dst_dir=image_dst_dir,
+                label_dst_dir=label_dst_dir,
+                generate_npy=generate_npy,
+                min_size=min_size,
+                max_size=max_size
+            )
+
+        if split == "train":
+            _generate_random_indices(len(image_src_paths), os.path.join(data_dst_dir, split))
+
+def _nwpu(
+    data_src_dir: str,
+    data_dst_dir: str,
+    min_size: int,
+    max_size: int,
+    generate_npy: bool = False
+) -> None:
+    for split in ["train", "val"]:
+        generate_npy = generate_npy and split == "train"
+        print(f"Processing {split}...")
+        with open(os.path.join(data_src_dir, f"{split}.txt"), "r") as f:
+            indices = f.read().splitlines()
+        indices = [idx.split(" ")[0] for idx in indices]
+        image_src_paths = [os.path.join(data_src_dir, f"images_part{min(5, (int(idx) - 1) // 1000 + 1)}", f"{idx}.jpg") for idx in indices]
+        label_src_paths = [os.path.join(data_src_dir, "mats", f"{idx}.mat") for idx in indices]
+
+        image_dst_dir = os.path.join(data_dst_dir, split, "images")
+        label_dst_dir = os.path.join(data_dst_dir, split, "labels")
+        os.makedirs(image_dst_dir, exist_ok=True)
+        os.makedirs(label_dst_dir, exist_ok=True)
+
+        size = len(str(len(image_src_paths)))
+        for i, (image_src_path, label_src_path) in tqdm(enumerate(zip(image_src_paths, label_src_paths)), total=len(image_src_paths)):
+            image_id = os.path.basename(image_src_path).split(".")[0]
+            label_id = os.path.basename(label_src_path).split(".")[0]
+            assert image_id == label_id, f"Expected image id {image_id} to match label id {label_id}"
+            name = f"{(i + 1):0{size}d}"
+            image = cv2.imread(image_src_path)
+            label = loadmat(label_src_path)["annPoints"]
+            _resize_and_save(
+                image=image,
+                label=label,
+                name=name,
+                image_dst_dir=image_dst_dir,
+                label_dst_dir=label_dst_dir,
+                generate_npy=generate_npy,
+                min_size=min_size,
+                max_size=max_size
+            )
+
+        if split == "train":
+            _generate_random_indices(len(image_src_paths), os.path.join(data_dst_dir, split))
+    
+    # preprocess the test set
+    split = "test"
+    print(f"Processing {split}...")
+    with open(os.path.join(data_src_dir, f"{split}.txt"), "r") as f:
+        indices = f.read().splitlines()
+    indices = [idx.split(" ")[0] for idx in indices]
+    image_src_paths = [os.path.join(data_src_dir, f"images_part{min(5, (int(idx) - 1) // 1000 + 1)}", f"{idx}.jpg") for idx in indices]
+
+    image_dst_dir = os.path.join(data_dst_dir, split, "images")
+    os.makedirs(image_dst_dir, exist_ok=True)
+
+    for image_src_path in tqdm(image_src_paths):
+        image_id = os.path.basename(image_src_path).split(".")[0]
+        image = cv2.imread(image_src_path)
+        _resize_and_save(
+            image=image,
+            label=None,
+            name=image_id,
+            image_dst_dir=image_dst_dir,
+            label_dst_dir=None,
+            generate_npy=generate_npy,
+            min_size=min_size,
+            max_size=max_size
+        )
+
+
+def _qnrf(
+    data_src_dir: str,
+    data_dst_dir: str,
+    min_size: int,
+    max_size: int,
+    generate_npy: bool = False
+) -> None:
+    for split in ["train", "val"]:
+        generate_npy = generate_npy and split == "train"
+        print(f"Processing {split}...")
+        if split == "train":
+            image_src_dir = os.path.join(data_src_dir, "Train")
+            label_src_dir = os.path.join(data_src_dir, "Train")
+            image_src_paths = glob(os.path.join(image_src_dir, "*.jpg"))
+            label_src_paths = glob(os.path.join(label_src_dir, "*.mat"))
+            assert len(image_src_paths) == len(label_src_paths) == 1201, f"Expected 1201 images and labels, got {len(image_src_paths)} images and {len(label_src_paths)} labels"
+        else:
+            image_src_dir = os.path.join(data_src_dir, "Test")
+            label_src_dir = os.path.join(data_src_dir, "Test")
+            image_src_paths = glob(os.path.join(image_src_dir, "*.jpg"))
+            label_src_paths = glob(os.path.join(label_src_dir, "*.mat"))
+            assert len(image_src_paths) == len(label_src_paths) == 334, f"Expected 334 images and labels, got {len(image_src_paths)} images and {len(label_src_paths)} labels"
+        
+        sort_key = lambda x: int((os.path.basename(x).split(".")[0]).split("_")[1])
+        image_src_paths.sort(key=sort_key)
+        label_src_paths.sort(key=sort_key)
+
+        image_dst_dir = os.path.join(data_dst_dir, split, "images")
+        label_dst_dir = os.path.join(data_dst_dir, split, "labels")
+        os.makedirs(image_dst_dir, exist_ok=True)
+        os.makedirs(label_dst_dir, exist_ok=True)
+    
+        size = len(str(len(image_src_paths)))
+        for i, (image_src_path, label_src_path) in tqdm(enumerate(zip(image_src_paths, label_src_paths)), total=len(image_src_paths)):
+            image_id = int((os.path.basename(image_src_path).split(".")[0]).split("_")[1])
+            label_id = int((os.path.basename(label_src_path).split(".")[0]).split("_")[1])
+            assert image_id == label_id, f"Expected image id {image_id} to match label id {label_id}"
+            name = f"{(i + 1):0{size}d}"
+            image = cv2.imread(image_src_path)
+            label = loadmat(label_src_path)["annPoints"]
+            _resize_and_save(
+                image=image,
+                label=label,
+                name=name,
+                image_dst_dir=image_dst_dir,
+                label_dst_dir=label_dst_dir,
+                generate_npy=generate_npy,
+                min_size=min_size,
+                max_size=max_size
+            )
+
+        if split == "train":
+            _generate_random_indices(len(image_src_paths), os.path.join(data_dst_dir, split))
+
+def _jhu(
+    data_src_dir: str,
+    data_dst_dir: str,
+    min_size: int,
+    max_size: int,
+    generate_npy: bool = False
+) -> None:
+    for split in ["train", "val"]:
+        generate_npy = generate_npy and split == "train"
+        if split == "train":
+            with open(os.path.join(data_src_dir, "train", "image_labels.txt"), "r") as f:
+                train_names = f.read().splitlines()
+            train_names = [name.split(",")[0] for name in train_names]
+            train_image_src_paths = [os.path.join(data_src_dir, "train", "images", f"{name}.jpg") for name in train_names]
+            train_label_src_paths = [os.path.join(data_src_dir, "train", "gt", f"{name}.txt") for name in train_names]
+
+            with open(os.path.join(data_src_dir, "val", "image_labels.txt"), "r") as f:
+                val_names = f.read().splitlines()
+            val_names = [name.split(",")[0] for name in val_names]
+            val_image_src_paths = [os.path.join(data_src_dir, "val", "images", f"{name}.jpg") for name in val_names]
+            val_label_src_paths = [os.path.join(data_src_dir, "val", "gt", f"{name}.txt") for name in val_names]
+
+            image_src_paths = train_image_src_paths + val_image_src_paths
+            label_src_paths = train_label_src_paths + val_label_src_paths
+
+        else:
+            with open(os.path.join(data_src_dir, "test", "image_labels.txt"), "r") as f:
+                test_names = f.read().splitlines()
+            test_names = [name.split(",")[0] for name in test_names]
+            image_src_paths = [os.path.join(data_src_dir, "test", "images", f"{name}.jpg") for name in test_names]
+            label_src_paths = [os.path.join(data_src_dir, "test", "gt", f"{name}.txt") for name in test_names]
+
+        image_dst_dir = os.path.join(data_dst_dir, split, "images")
+        label_dst_dir = os.path.join(data_dst_dir, split, "labels")
+        os.makedirs(image_dst_dir, exist_ok=True)
+        os.makedirs(label_dst_dir, exist_ok=True)
+
+        size = len(str(len(image_src_paths)))
+        for i, (image_src_path, label_src_path) in tqdm(enumerate(zip(image_src_paths, label_src_paths)), total=len(image_src_paths)):
+            image_id = int(os.path.basename(image_src_path).split(".")[0])
+            label_id = int(os.path.basename(label_src_path).split(".")[0])
+            assert image_id == label_id, f"Expected image id {image_id} to match label id {label_id}"
+            name = f"{(i + 1):0{size}d}"
+            image = cv2.imread(image_src_path)
+            with open(label_src_path, "r") as f:
+                label = f.read().splitlines()
+            label = np.array([list(map(float, line.split(" ")[0: 2])) for line in label])
+            _resize_and_save(
+                image=image,
+                label=label,
+                name=name,
+                image_dst_dir=image_dst_dir,
+                label_dst_dir=label_dst_dir,
+                generate_npy=generate_npy,
+                min_size=min_size,
+                max_size=max_size
+            )
+
+        if split == "train":
+            _generate_random_indices(len(image_src_paths), os.path.join(data_dst_dir, split))
+
+
+def parse_args():
+    parser = ArgumentParser(description="Pre-process datasets to resize images and labeld into a given range.")
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        choices=["nwpu", "ucf_qnrf", "jhu", "shanghaitech_a", "shanghaitech_b"],
+        required=True,
+        help="The dataset to pre-process."
+    )
+    parser.add_argument("--src_dir", type=str, required=True, help="The root directory of the source dataset.")
+    parser.add_argument("--dst_dir", type=str, required=True, help="The root directory of the destination dataset.")
+    parser.add_argument("--min_size", type=int, default=256, help="The minimum size of the shorter side of the image.")
+    parser.add_argument("--max_size", type=int, default=None, help="The maximum size of the longer side of the image.")
+    parser.add_argument("--generate_npy", action="store_true", help="Generate .npy files for images.")
+
+    args = parser.parse_args()
+    args.src_dir = os.path.abspath(args.src_dir)
+    args.dst_dir = os.path.abspath(args.dst_dir)
+    args.max_size = float("inf") if args.max_size is None else args.max_size
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    _preprocess(
+        dataset=args.dataset,
+        data_src_dir=args.src_dir,
+        data_dst_dir=args.dst_dir,
+        min_size=args.min_size,
+        max_size=args.max_size,
+        generate_npy=args.generate_npy
+    )

preprocess.sh

@@ -0,0 +1,8 @@
+#!/bin/sh
+# !!! ----------------- Please change the `src_dir` to the location of your dataset. ----------------- !!!
+# !!! ----- Please do NOT change `dst_dir` as it will be used to locate the dataset in crowd.py. ----- !!!
+
+python preprocess.py --dataset shanghaitech_a --src_dir ./data/ShanghaiTech/part_A --dst_dir ./data/sha  --min_size 448 --max_size 4096
+python preprocess.py --dataset shanghaitech_b --src_dir ./data/ShanghaiTech/part_B --dst_dir ./data/shb  --min_size 448 --max_size 4096
+python preprocess.py --dataset nwpu           --src_dir ./data/NWPU-Crowd          --dst_dir ./data/nwpu --min_size 448 --max_size 3072
+python preprocess.py --dataset ucf_qnrf       --src_dir ./data/UCF-QNRF            --dst_dir ./data/qnrf --min_size 448 --max_size 2048

requirements.txt

@@ -0,0 +1,13 @@
+einops==0.7.0
+ftfy==6.1.3
+numpy==1.26.4
+Pillow==10.2.0
+regex==2023.12.25
+scipy==1.12.0
+setuptools==69.1.1
+tensorboardX==2.6.2.2
+tensorboardX==2.6.2.2
+timm==0.9.16
+torch==2.2.1
+torchvision==0.17.1
+tqdm==4.66.2

run.sh

@@ -0,0 +1,32 @@
+#!/bin/sh
+export CUDA_VISIBLE_DEVICES=0  # Set the GPU ID. Comment this line to use all GPUs and DDP.
+
+# Train the commonly used VGG19-based encoder-decoder model on NWPU-Crowd.
+# Change `--dataset` to `sha` or `shb` or `qnrf` to train on ShanghaiTech A, or ShanghaiTech B, or UCF-QNRF.
+python trainer.py \
+    --model vgg19_ae --input_size 448 --reduction 8 --truncation 4 --anchor_points average \
+    --dataset nwpu \
+    --count_loss dmcount &&
+
+# Train the CLIP-EBC (ResNet50) model on ShanghaiTech A. Use `--dataset shb` if you want to train on ShanghaiTech B.
+python trainer.py \
+    --model clip_resnet50 --input_size 448 --reduction 8 --truncation 4 --anchor_points average --prompt_type word \
+    --dataset sha \
+    # --sliding_window --window_size 448 --stride 448 \  # Uncomment this line to enable sliding window prediction with a stride size of 448.
+    --count_loss dmcount &&
+
+# Train the CLIP-EBC (ViT-B/16) model on UCF-QNRF, using VPT in training and sliding window prediction in testing.
+# By default, 32 tokens for each layer are used in VPT. You can also set `--num_vpt` to change the number of tokens.
+# By default, the deep visual prompt tuning is used. You can set `--shallow_vpt` to use the shallow visual prompt tuning.
+# `--amp` enables automatic mixed precision training.
+python trainer.py \
+    --model clip_vit_b_16 --input_size 224 --reduction 8 --truncation 4 \
+    --dataset qnrf --batch_size 16 --amp \
+    --num_crops 2 --sliding_window --window_size 224 --stride 224 --warmup_lr 1e-3 \
+    --count_loss dmcount
+
+# Generate results on NWPU-Crowd Test.
+# python test_nwpu.py \
+#     --model clip_vit_b_16 --input_size 224 --reduction 8 --truncation 4 --anchor_points average --prompt_type word \
+#     --num_vpt 32 --vpt_drop 0.0 --sliding_window --stride 224 \
+#     --weight_path ./checkpoints/nwpu/clip_vit_b_16_word_224_8_4_fine_1.0_dmcount/best_rmse_1.pth

test_nwpu.py

@@ -0,0 +1,152 @@
+import torch
+from argparse import ArgumentParser
+import os, json
+from tqdm import tqdm
+
+current_dir = os.path.abspath(os.path.dirname(__file__))
+
+from datasets import NWPUTest, Resize2Multiple
+from models import get_model
+from utils import get_config, sliding_window_predict
+
+parser = ArgumentParser(description="Test a trained model on the NWPU-Crowd test set.")
+# Parameters for model
+parser.add_argument("--model", type=str, default="vgg19_ae", help="The model to train.")
+parser.add_argument("--input_size", type=int, default=448, help="The size of the input image.")
+parser.add_argument("--reduction", type=int, default=8, choices=[8, 16, 32], help="The reduction factor of the model.")
+parser.add_argument("--regression", action="store_true", help="Use blockwise regression instead of classification.")
+parser.add_argument("--truncation", type=int, default=None, help="The truncation of the count.")
+parser.add_argument("--anchor_points", type=str, default="average", choices=["average", "middle"], help="The representative count values of bins.")
+parser.add_argument("--prompt_type", type=str, default="word", choices=["word", "number"], help="The prompt type for CLIP.")
+parser.add_argument("--granularity", type=str, default="fine", choices=["fine", "dynamic", "coarse"], help="The granularity of bins.")
+parser.add_argument("--num_vpt", type=int, default=32, help="The number of visual prompt tokens.")
+parser.add_argument("--vpt_drop", type=float, default=0.0, help="The dropout rate for visual prompt tokens.")
+parser.add_argument("--shallow_vpt", action="store_true", help="Use shallow visual prompt tokens.")
+parser.add_argument("--weight_path", type=str, required=True, help="The path to the weights of the model.")
+
+# Parameters for evaluation
+parser.add_argument("--sliding_window", action="store_true", help="Use sliding window strategy for evaluation.")
+parser.add_argument("--stride", type=int, default=None, help="The stride for sliding window strategy.")
+parser.add_argument("--window_size", type=int, default=None, help="The window size for in prediction.")
+parser.add_argument("--resize_to_multiple", action="store_true", help="Resize the image to the nearest multiple of the input size.")
+parser.add_argument("--zero_pad_to_multiple", action="store_true", help="Zero pad the image to the nearest multiple of the input size.")
+
+parser.add_argument("--device", type=str, default="cuda", help="The device to use for evaluation.")
+parser.add_argument("--num_workers", type=int, default=4, help="The number of workers for the data loader.")
+
+
+def main(args: ArgumentParser):
+    print("Testing a trained model on the NWPU-Crowd test set.")
+    device = torch.device(args.device)
+    _ = get_config(vars(args).copy(), mute=False)
+    if args.regression:
+        bins, anchor_points = None, None
+    else:
+        with open(os.path.join(current_dir, "configs", f"reduction_{args.reduction}.json"), "r") as f:
+            config = json.load(f)[str(args.truncation)]["nwpu"]
+        bins = config["bins"][args.granularity]
+        anchor_points = config["anchor_points"][args.granularity]["average"] if args.anchor_points == "average" else config["anchor_points"][args.granularity]["middle"]
+        bins = [(float(b[0]), float(b[1])) for b in bins]
+        anchor_points = [float(p) for p in anchor_points]
+
+    args.bins = bins
+    args.anchor_points = anchor_points
+
+    model = get_model(
+        backbone=args.model,
+        input_size=args.input_size, 
+        reduction=args.reduction,
+        bins=bins,
+        anchor_points=anchor_points,
+        prompt_type=args.prompt_type,
+        num_vpt=args.num_vpt,
+        vpt_drop=args.vpt_drop,
+        deep_vpt=not args.shallow_vpt
+    )
+    state_dict = torch.load(args.weight_path, map_location="cpu")
+    state_dict = state_dict if "best" in os.path.basename(args.weight_path) else state_dict["model_state_dict"]
+    model.load_state_dict(state_dict, strict=True)
+    model = model.to(device)
+    model.eval()
+
+    sliding_window = args.sliding_window
+    if args.sliding_window:
+        window_size = args.input_size
+        stride = window_size // 2 if args.stride is None else args.stride
+        if args.resize_to_multiple:
+            transforms = Resize2Multiple(base=args.input_size)
+        else:
+            transforms = None
+    else:
+        window_size, stride = None, None
+        transforms = None
+
+    dataset = NWPUTest(transforms=transforms, return_filename=True)
+
+    image_ids = []
+    preds = []
+
+    for idx in tqdm(range(len(dataset)), desc="Testing on NWPU"):
+        image, image_path = dataset[idx]
+        image = image.unsqueeze(0)  # add batch dimension
+        image = image.to(device)  # add batch dimension
+
+        with torch.set_grad_enabled(False):
+            if sliding_window:
+                pred_density = sliding_window_predict(model, image, window_size, stride)
+            else:
+                pred_density = model(image)
+
+            pred_count = pred_density.sum(dim=(1, 2, 3)).item()
+
+        image_ids.append(os.path.basename(image_path).split(".")[0])
+        preds.append(pred_count)
+
+    result_dir = os.path.join(current_dir, "nwpu_test_results")
+    os.makedirs(result_dir, exist_ok=True)
+    weights_dir, weights_name = os.path.split(args.weight_path)
+    model_name = os.path.split(weights_dir)[-1]
+    result_path = os.path.join(result_dir, f"{model_name}_{weights_name.split('.')[0]}.txt")
+
+    with open(result_path, "w") as f:
+        for idx, (image_id, pred) in enumerate(zip(image_ids, preds)):
+            if idx != len(image_ids) - 1:
+                f.write(f"{image_id} {pred}\n")
+            else:
+                f.write(f"{image_id} {pred}")  # no newline at the end of the file
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    args.model = args.model.lower()
+
+    if args.regression:
+        args.truncation = None
+        args.anchor_points = None
+        args.bins = None
+        args.prompt_type = None
+        args.granularity = None
+
+    if "clip_vit" not in args.model:
+        args.num_vpt = None
+        args.vpt_drop = None
+        args.shallow_vpt = None
+    
+    if "clip" not in args.model:
+        args.prompt_type = None
+
+    if args.sliding_window:
+        args.window_size = args.input_size if args.window_size is None else args.window_size
+        args.stride = args.input_size if args.stride is None else args.stride
+        assert not (args.zero_pad_to_multiple and args.resize_to_multiple), "Cannot use both zero pad and resize to multiple."
+
+    else:
+        args.window_size = None
+        args.stride = None
+        args.zero_pad_to_multiple = False
+        args.resize_to_multiple = False
+    
+    main(args)
+
+# Example usage:
+# python test_nwpu.py --model vgg19_ae --truncation 4 --weight_path ./checkpoints/sha/vgg19_ae_448_4_1.0_dmcount_aug/best_mae.pth --device cuda:0

test_nwpu.sh

@@ -0,0 +1,7 @@
+#!/bin/sh
+export CUDA_VISIBLE_DEVICES=0
+
+python test_nwpu.py \
+    --model clip_vit_b_16 --input_size 224 --reduction 8 --truncation 4 --anchor_points average --prompt_type word \
+    --num_vpt 32 --vpt_drop 0.0 --sliding_window --stride 224 \
+    --weight_path ./checkpoints/nwpu/clip_vit_b_16_word_224_8_4_fine_1.0_dmcount/best_mae_0.pth

train.py

@@ -0,0 +1,69 @@
+import torch
+from torch import nn
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader
+from torch.cuda.amp import GradScaler, autocast
+import numpy as np
+from tqdm import tqdm
+from typing import Dict, Tuple
+
+
+from utils import barrier, reduce_mean, update_loss_info
+
+
+def train(
+    model: nn.Module,
+    data_loader: DataLoader,
+    loss_fn: nn.Module,
+    optimizer: Optimizer,
+    grad_scaler: GradScaler,
+    device: torch.device,
+    rank: int,
+    nprocs: int,
+) -> Tuple[nn.Module, Optimizer, GradScaler, Dict[str, float]]:
+    model.train()
+    info = None
+    data_iter = tqdm(data_loader) if rank == 0 else data_loader
+    ddp = nprocs > 1
+    regression = (model.module.bins is None) if ddp else (model.bins is None)
+
+    for image, target_points, target_density in data_iter:
+        image = image.to(device)
+        target_points = [p.to(device) for p in target_points]
+        target_density = target_density.to(device)
+        with torch.set_grad_enabled(True):
+
+            if grad_scaler is not None:
+                with autocast(enabled=grad_scaler.is_enabled()):
+                    if not regression:
+                        pred_class, pred_density = model(image)
+                        loss, loss_info = loss_fn(pred_class, pred_density, target_density, target_points)
+                    else:
+                        pred_density = model(image)
+                        loss, loss_info = loss_fn(pred_density, target_density, target_points)
+
+            else:
+                if not regression:
+                    pred_class, pred_density = model(image)
+                    loss, loss_info = loss_fn(pred_class, pred_density, target_density, target_points)
+                else:
+                    pred_density = model(image)
+                    loss, loss_info = loss_fn(pred_density, target_density, target_points)
+
+        optimizer.zero_grad()
+        if grad_scaler is not None:
+            grad_scaler.scale(loss).backward()
+            grad_scaler.step(optimizer)
+            grad_scaler.update()
+        else:
+            loss.backward()
+            optimizer.step()
+
+        loss_info = {k: reduce_mean(v.detach(), nprocs).item() if ddp else v.detach().item() for k, v in loss_info.items()}
+        # if rank == 0:
+            # loss_info = {k: v.item() for k, v in loss_info.items()}
+        info = update_loss_info(info, loss_info)
+
+        barrier(ddp)
+
+    return model, optimizer, grad_scaler, {k: np.mean(v) for k, v in info.items()}

trainer.py

@@ -0,0 +1,246 @@
+import torch
+from torch import nn
+import torch.multiprocessing as mp
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.cuda.amp import GradScaler
+
+from argparse import ArgumentParser
+import os, json
+
+current_dir = os.path.abspath(os.path.dirname(__file__))
+
+from datasets import standardize_dataset_name
+from models import get_model
+
+from utils import setup, cleanup, init_seeds, get_logger, get_config, barrier
+from utils import get_dataloader, get_loss_fn, get_optimizer, load_checkpoint, save_checkpoint
+from utils import get_writer, update_train_result, update_eval_result, log
+from train import train
+from eval import evaluate
+
+
+parser = ArgumentParser(description="Train an EBC model.")
+
+# Parameters for model
+parser.add_argument("--model", type=str, default="vgg19_ae", help="The model to train.")
+parser.add_argument("--input_size", type=int, default=448, help="The size of the input image.")
+parser.add_argument("--reduction", type=int, default=8, choices=[8, 16, 32], help="The reduction factor of the model.")
+parser.add_argument("--regression", action="store_true", help="Use blockwise regression instead of classification.")
+parser.add_argument("--truncation", type=int, default=None, help="The truncation of the count.")
+parser.add_argument("--anchor_points", type=str, default="average", choices=["average", "middle"], help="The representative count values of bins.")
+parser.add_argument("--prompt_type", type=str, default="word", choices=["word", "number"], help="The prompt type for CLIP.")
+parser.add_argument("--granularity", type=str, default="fine", choices=["fine", "dynamic", "coarse"], help="The granularity of bins.")
+parser.add_argument("--num_vpt", type=int, default=32, help="The number of visual prompt tokens.")
+parser.add_argument("--vpt_drop", type=float, default=0.0, help="The dropout rate for visual prompt tokens.")
+parser.add_argument("--shallow_vpt", action="store_true", help="Use shallow visual prompt tokens.")
+
+# Parameters for dataset
+parser.add_argument("--dataset", type=str, required=True, help="The dataset to train on.")
+parser.add_argument("--batch_size", type=int, default=8, help="The training batch size.")
+parser.add_argument("--num_crops", type=int, default=1, help="The number of crops for multi-crop training.")
+parser.add_argument("--min_scale", type=float, default=1.0, help="The minimum scale for random scale augmentation.")
+parser.add_argument("--max_scale", type=float, default=2.0, help="The maximum scale for random scale augmentation.")
+parser.add_argument("--brightness", type=float, default=0.1, help="The brightness factor for random color jitter augmentation.")
+parser.add_argument("--contrast", type=float, default=0.1, help="The contrast factor for random color jitter augmentation.")
+parser.add_argument("--saturation", type=float, default=0.1, help="The saturation factor for random color jitter augmentation.")
+parser.add_argument("--hue", type=float, default=0.0, help="The hue factor for random color jitter augmentation.")
+parser.add_argument("--kernel_size", type=int, default=5, help="The kernel size for Gaussian blur augmentation.")
+parser.add_argument("--saltiness", type=float, default=1e-3, help="The saltiness for pepper salt noise augmentation.")
+parser.add_argument("--spiciness", type=float, default=1e-3, help="The spiciness for pepper salt noise augmentation.")
+parser.add_argument("--jitter_prob", type=float, default=0.2, help="The probability for random color jitter augmentation.")
+parser.add_argument("--blur_prob", type=float, default=0.2, help="The probability for Gaussian blur augmentation.")
+parser.add_argument("--noise_prob", type=float, default=0.5, help="The probability for pepper salt noise augmentation.")
+
+# Parameters for evaluation
+parser.add_argument("--sliding_window", action="store_true", help="Use sliding window strategy for evaluation.")
+parser.add_argument("--stride", type=int, default=None, help="The stride for sliding window strategy.")
+parser.add_argument("--window_size", type=int, default=None, help="The window size for in prediction.")
+parser.add_argument("--resize_to_multiple", action="store_true", help="Resize the image to the nearest multiple of the input size.")
+parser.add_argument("--zero_pad_to_multiple", action="store_true", help="Zero pad the image to the nearest multiple of the input size.")
+
+# Parameters for loss function
+parser.add_argument("--weight_count_loss", type=float, default=1.0, help="The weight for count loss.")
+parser.add_argument("--count_loss", type=str, default="mae", choices=["mae", "mse", "dmcount"], help="The loss function for count.")
+
+# Parameters for optimizer (Adam)
+parser.add_argument("--lr", type=float, default=1e-4, help="The learning rate.")
+parser.add_argument("--weight_decay", type=float, default=1e-4, help="The weight decay.")
+
+# Parameters for learning rate scheduler
+parser.add_argument("--warmup_epochs", type=int, default=50, help="Number of epochs for warmup. The learning rate will increase from eta_min to lr.")
+parser.add_argument("--warmup_lr", type=float, default=1e-6, help="Learning rate for warmup.")
+parser.add_argument("--T_0", type=int, default=5, help="Number of epochs for the first restart.")
+parser.add_argument("--T_mult", type=int, default=2, help="A factor increases T_0 after a restart.")
+parser.add_argument("--eta_min", type=float, default=1e-7, help="Minimum learning rate.")
+
+# Parameters for training
+parser.add_argument("--total_epochs", type=int, default=2600, help="Number of epochs to train.")
+parser.add_argument("--eval_start", type=int, default=50, help="Start to evaluate after this number of epochs.")
+parser.add_argument("--eval_freq", type=int, default=1, help="Evaluate every this number of epochs.")
+parser.add_argument("--save_freq", type=int, default=5, help="Save checkpoint every this number of epochs. Could help reduce I/O.")
+parser.add_argument("--save_best_k", type=int, default=3, help="Save the best k checkpoints.")
+parser.add_argument("--amp", action="store_true", help="Use automatic mixed precision training.")
+parser.add_argument("--num_workers", type=int, default=4, help="Number of workers for data loading.")
+parser.add_argument("--local_rank", type=int, default=-1, help="Local rank for distributed training.")
+parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+
+
+def run(local_rank: int, nprocs: int, args: ArgumentParser) -> None:
+    print(f"Rank {local_rank} process among {nprocs} processes.")
+    init_seeds(args.seed + local_rank)
+    setup(local_rank, nprocs)
+    print(f"Initialized successfully. Training with {nprocs} GPUs.")
+    device = f"cuda:{local_rank}" if local_rank != -1 else "cuda:0"
+    print(f"Using device: {device}.")
+
+    ddp = nprocs > 1
+
+    if args.regression:
+        bins, anchor_points = None, None
+    else:
+        with open(os.path.join(current_dir, "configs", f"reduction_{args.reduction}.json"), "r") as f:
+            config = json.load(f)[str(args.truncation)][args.dataset]
+        bins = config["bins"][args.granularity]
+        anchor_points = config["anchor_points"][args.granularity]["average"] if args.anchor_points == "average" else config["anchor_points"][args.granularity]["middle"]
+        bins = [(float(b[0]), float(b[1])) for b in bins]
+        anchor_points = [float(p) for p in anchor_points]
+
+    args.bins = bins
+    args.anchor_points = anchor_points
+
+    model = get_model(
+        backbone=args.model,
+        input_size=args.input_size, 
+        reduction=args.reduction,
+        bins=bins,
+        anchor_points=anchor_points,
+        prompt_type=args.prompt_type,
+        num_vpt=args.num_vpt,
+        vpt_drop=args.vpt_drop,
+        deep_vpt=not args.shallow_vpt
+    ).to(device)
+
+    grad_scaler = GradScaler() if args.amp else None
+
+    loss_fn = get_loss_fn(args).to(device)
+    optimizer, scheduler = get_optimizer(args, model)
+
+    ckpt_dir_name = f"{args.model}_{args.prompt_type}_" if "clip" in args.model else f"{args.model}_"
+    ckpt_dir_name += f"{args.input_size}_{args.reduction}_{args.truncation}_{args.granularity}_"
+    ckpt_dir_name += f"{args.weight_count_loss}_{args.count_loss}"
+
+    args.ckpt_dir = os.path.join(current_dir, "checkpoints", args.dataset, ckpt_dir_name)
+    os.makedirs(args.ckpt_dir, exist_ok=True)
+    model, optimizer, scheduler, grad_scaler, start_epoch, loss_info, hist_val_scores, best_val_scores = load_checkpoint(args, model, optimizer, scheduler, grad_scaler)
+
+    if local_rank == 0:
+        model_without_ddp = model
+        writer = get_writer(args.ckpt_dir)
+        logger = get_logger(os.path.join(args.ckpt_dir, "train.log"))
+        logger.info(get_config(vars(args), mute=False))
+        val_loader = get_dataloader(args, split="val", ddp=False)
+
+    args.batch_size = int(args.batch_size / nprocs)
+    args.num_workers = int(args.num_workers / nprocs)
+    train_loader, sampler = get_dataloader(args, split="train", ddp=ddp)
+
+    model = DDP(nn.SyncBatchNorm.convert_sync_batchnorm(model), device_ids=[local_rank], output_device=local_rank) if ddp else model
+
+    for epoch in range(start_epoch, args.total_epochs + 1):  # start from 1
+        if local_rank == 0:
+            message = f"\tlr: {optimizer.param_groups[0]['lr']:.3e}"
+            log(logger, epoch, args.total_epochs, message=message)
+
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+
+        model, optimizer, grad_scaler, loss_info = train(model, train_loader, loss_fn, optimizer, grad_scaler, device, local_rank, nprocs)
+        scheduler.step()
+        barrier(ddp)
+
+        if local_rank == 0:
+            eval = (epoch >= args.eval_start) and ((epoch - args.eval_start) % args.eval_freq == 0)
+            update_train_result(epoch, loss_info, writer)
+            log(logger, None, None, loss_info=loss_info, message="\n" * 2 if not eval else None)
+
+            if eval:
+                print("Evaluating")
+                state_dict = model.module.state_dict() if ddp else model.state_dict()
+                model_without_ddp.load_state_dict(state_dict)
+                curr_val_scores = evaluate(
+                    model_without_ddp,
+                    val_loader,
+                    device,
+                    args.sliding_window,
+                    args.input_size,
+                    args.stride,
+                )
+                hist_val_scores, best_val_scores = update_eval_result(epoch, curr_val_scores, hist_val_scores, best_val_scores, writer, state_dict, os.path.join(args.ckpt_dir))
+                log(logger, None, None, None, curr_val_scores, best_val_scores, message="\n" * 3)
+    
+            if (epoch % args.save_freq == 0):
+                save_checkpoint(
+                    epoch + 1,
+                    model.module.state_dict() if ddp else model.state_dict(),
+                    optimizer.state_dict(),
+                    scheduler.state_dict() if scheduler is not None else None,
+                    grad_scaler.state_dict() if grad_scaler is not None else None,
+                    loss_info,
+                    hist_val_scores,
+                    best_val_scores,
+                    args.ckpt_dir,
+                )
+
+        barrier(ddp)
+
+    if local_rank == 0:
+        writer.close()
+        print("Training completed. Best scores:")
+        for k in best_val_scores.keys():
+            scores = " ".join([f"{best_val_scores[k][i]:.4f};" for i in range(len(best_val_scores[k]))])
+            print(f"    {k}: {scores}")
+
+    cleanup(ddp)
+
+
+def main():
+    args = parser.parse_args()
+    args.model = args.model.lower()
+    args.dataset = standardize_dataset_name(args.dataset)
+
+    if args.regression:
+        args.truncation = None
+        args.anchor_points = None
+        args.bins = None
+        args.prompt_type = None
+        args.granularity = None
+
+    if "clip_vit" not in args.model:
+        args.num_vpt = None
+        args.vpt_drop = None
+        args.shallow_vpt = None
+    
+    if "clip" not in args.model:
+        args.prompt_type = None
+
+    if args.sliding_window:
+        args.window_size = args.input_size if args.window_size is None else args.window_size
+        args.stride = args.input_size if args.stride is None else args.stride
+        assert not (args.zero_pad_to_multiple and args.resize_to_multiple), "Cannot use both zero pad and resize to multiple."
+
+    else:
+        args.window_size = None
+        args.stride = None
+        args.zero_pad_to_multiple = False
+        args.resize_to_multiple = False
+
+    args.nprocs = torch.cuda.device_count()
+    print(f"Using {args.nprocs} GPUs.")
+    if args.nprocs > 1:
+        mp.spawn(run, nprocs=args.nprocs, args=(args.nprocs, args))
+    else:
+        run(0, 1, args)
+
+
+if __name__ == "__main__":
+    main()

utils/__init__.py

@@ -0,0 +1,13 @@
+from .ddp_utils import reduce_mean, setup, cleanup, init_seeds, barrier
+from .eval_utils import calculate_errors, resize_density_map, sliding_window_predict
+from .log_utils import get_logger, get_config, get_writer, print_epoch, print_train_result, print_eval_result, update_train_result, update_eval_result, log, update_loss_info
+from .train_utils import cosine_annealing_warm_restarts, get_loss_fn, get_optimizer, load_checkpoint, save_checkpoint
+from .data_utils import get_dataloader
+
+
+__all__ = [
+    "reduce_mean", "setup", "cleanup", "init_seeds", "barrier",
+    "calculate_errors", "resize_density_map", "sliding_window_predict",
+    "get_logger", "get_config", "get_writer", "print_epoch", "print_train_result", "print_eval_result", "update_train_result", "update_eval_result", "log", "update_loss_info",
+    "get_dataloader", "get_loss_fn", "get_optimizer", "load_checkpoint", "save_checkpoint",
+]

utils/data_utils.py

@@ -0,0 +1,78 @@
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torchvision.transforms.v2 import Compose
+import os, sys
+from argparse import ArgumentParser
+from typing import Union, Tuple
+
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(parent_dir)
+
+import datasets
+
+
+def get_dataloader(args: ArgumentParser, split: str = "train", ddp: bool = False) -> Union[Tuple[DataLoader, Union[DistributedSampler, None]], DataLoader]:
+    if split == "train":  # train, strong augmentation
+        transforms = Compose([
+            datasets.RandomResizedCrop((args.input_size, args.input_size), scale=(args.min_scale, args.max_scale)),
+            datasets.RandomHorizontalFlip(),
+            datasets.RandomApply([
+                datasets.ColorJitter(brightness=args.brightness, contrast=args.contrast, saturation=args.saturation, hue=args.hue),
+                datasets.GaussianBlur(kernel_size=args.kernel_size, sigma=(0.1, 5.0)),
+                datasets.PepperSaltNoise(saltiness=args.saltiness, spiciness=args.spiciness),
+            ], p=(args.jitter_prob, args.blur_prob, args.noise_prob)),
+        ])
+
+    elif args.sliding_window:
+        if args.resize_to_multiple:
+            transforms = datasets.Resize2Multiple(args.window_size, stride=args.stride)
+        elif args.zero_pad_to_multiple:
+            transforms = datasets.ZeroPad2Multiple(args.window_size, stride=args.stride)
+        else:
+            transforms = None
+
+    else:
+        transforms = None
+
+    dataset = datasets.Crowd(
+        dataset=args.dataset,
+        split=split,
+        transforms=transforms,
+        sigma=None,
+        return_filename=False,
+        num_crops=args.num_crops if split == "train" else 1,
+    )
+
+    if ddp and split == "train":  # data_loader for training in DDP
+        sampler = DistributedSampler(dataset)
+        data_loader = DataLoader(
+            dataset,
+            batch_size=args.batch_size,
+            sampler=sampler,
+            num_workers=args.num_workers,
+            pin_memory=True,
+            collate_fn=datasets.collate_fn,
+        )
+        return data_loader, sampler
+
+    elif split == "train":  # data_loader for training
+        data_loader = DataLoader(
+            dataset,
+            batch_size=args.batch_size,
+            shuffle=True,
+            num_workers=args.num_workers,
+            pin_memory=True,
+            collate_fn=datasets.collate_fn,
+        )
+        return data_loader, None
+
+    else:  # data_loader for evaluation
+        data_loader = DataLoader(
+            dataset,
+            batch_size=1,  # Use batch size 1 for evaluation
+            shuffle=False,
+            num_workers=args.num_workers,
+            pin_memory=True,
+            collate_fn=datasets.collate_fn,
+        )
+        return data_loader

utils/ddp_utils.py

@@ -0,0 +1,44 @@
+import torch
+from torch import Tensor
+import torch.distributed as dist
+import numpy as np
+import random
+import os
+
+
+def reduce_mean(tensor: Tensor, nprocs: int) -> Tensor:
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    rt /= nprocs
+    return rt
+
+
+def setup(local_rank: int, nprocs: int) -> None:
+    if nprocs > 1:
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = "12355"
+        dist.init_process_group("nccl", rank=local_rank, world_size=nprocs)
+    else:
+        print("Single process. No need to setup dist.")
+
+
+def cleanup(ddp: bool = True) -> None:
+    if ddp:
+        dist.destroy_process_group()
+
+
+def init_seeds(seed: int, cuda_deterministic: bool = False) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if cuda_deterministic:  # slower, but reproducible
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    else:  # faster, not reproducible
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.benchmark = True
+
+
+def barrier(ddp: bool = True) -> None:
+    if ddp:
+        dist.barrier()