| import torch |
| from torch import nn |
| import numpy as np |
| import os, json |
| from tqdm import tqdm |
| from argparse import ArgumentParser |
| from typing import Dict |
|
|
| import datasets |
|
|
|
|
| class SumPool2d(nn.Module): |
| def __init__(self, kernel_size: int, stride: int): |
| super(SumPool2d, self).__init__() |
| self.kernel_size = kernel_size |
| self.stride = stride |
| self.sum_pool = nn.AvgPool2d(kernel_size, stride, divisor_override=1) |
|
|
| def forward(self, x): |
| return self.sum_pool(x) |
|
|
|
|
| def _update_dict(d: Dict, keys: np.ndarray, values: np.ndarray) -> Dict: |
| keys = keys.tolist() if isinstance(keys, np.ndarray) else keys |
| values = values.tolist() if isinstance(values, np.ndarray) else values |
| for k, v in zip(keys, values): |
| d[k] = d.get(k, 0) + v |
|
|
| return d |
|
|
|
|
| def _get_counts( |
| dataset_name: str, |
| device: torch.device, |
| ) -> None: |
| filter_4 = SumPool2d(4, 1).to(device) |
| filter_7 = SumPool2d(7, 1).to(device) |
| filter_8 = SumPool2d(8, 1).to(device) |
| filter_14 = SumPool2d(14, 1).to(device) |
| filter_16 = SumPool2d(16, 1).to(device) |
| filter_28 = SumPool2d(28, 1).to(device) |
| filter_32 = SumPool2d(32, 1).to(device) |
| filter_56 = SumPool2d(56, 1).to(device) |
| filter_64 = SumPool2d(64, 1).to(device) |
| counts_1, counts_4, counts_7, counts_8 = {}, {}, {}, {} |
| counts_14, counts_16 = {}, {} |
| counts_28, counts_32 = {}, {} |
| counts_56, counts_64 = {}, {} |
|
|
| max_counts_4 = {"max": 0., "name": None, "x": None, "y": None} |
| max_counts_7 = {"max": 0., "name": None, "x": None, "y": None} |
| max_counts_8 = {"max": 0., "name": None, "x": None, "y": None} |
| max_counts_14 = {"max": 0., "name": None, "x": None, "y": None} |
| max_counts_16 = {"max": 0., "name": None, "x": None, "y": None} |
| max_counts_28 = {"max": 0., "name": None, "x": None, "y": None} |
| max_counts_32 = {"max": 0., "name": None, "x": None, "y": None} |
| max_counts_56 = {"max": 0., "name": None, "x": None, "y": None} |
| max_counts_64 = {"max": 0., "name": None, "x": None, "y": None} |
|
|
| counts_dir = os.path.join(os.getcwd(), "counts") |
| os.makedirs(counts_dir, exist_ok=True) |
|
|
| dataset = datasets.Crowd(dataset=dataset_name, split="train", transforms=None, return_filename=True) |
| print(f"Counting {dataset_name} dataset") |
|
|
| for i in tqdm(range(len(dataset))): |
| _, _, density, img_name = dataset[i] |
| density_np = density.cpu().numpy().astype(int) |
| uniques_, counts_ = np.unique(density_np, return_counts=True) |
| counts_1 = _update_dict(counts_1, uniques_, counts_) |
|
|
| density = density.to(device) |
| window_4, window_7, window_8 = filter_4(density), filter_7(density), filter_8(density) |
| window_14, window_16 = filter_14(density), filter_16(density) |
| window_28, window_32 = filter_28(density), filter_32(density) |
| window_56, window_64 = filter_56(density), filter_64(density) |
|
|
| window_4, window_7, window_8 = torch.round(window_4).int(), torch.round(window_7).int(), torch.round(window_8).int() |
| window_14, window_16 = torch.round(window_14).int(), torch.round(window_16).int() |
| window_28, window_32 = torch.round(window_28).int(), torch.round(window_32).int() |
| window_56, window_64 = torch.round(window_56).int(), torch.round(window_64).int() |
|
|
| window_4, window_7, window_8 = torch.squeeze(window_4), torch.squeeze(window_7), torch.squeeze(window_8) |
| window_14, window_16 = torch.squeeze(window_14), torch.squeeze(window_16) |
| window_28, window_32 = torch.squeeze(window_28), torch.squeeze(window_32) |
| window_56, window_64 = torch.squeeze(window_56), torch.squeeze(window_64) |
|
|
| if window_4.max().item() > max_counts_4["max"]: |
| max_counts_4["max"] = window_4.max().item() |
| max_counts_4["name"] = img_name |
| x, y = torch.where(window_4 == window_4.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_4["x"] = x |
| max_counts_4["y"] = y |
| |
| if window_7.max().item() > max_counts_7["max"]: |
| max_counts_7["max"] = window_7.max().item() |
| max_counts_7["name"] = img_name |
| x, y = torch.where(window_7 == window_7.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_7["x"] = x |
| max_counts_7["y"] = y |
| |
| if window_8.max().item() > max_counts_8["max"]: |
| max_counts_8["max"] = window_8.max().item() |
| max_counts_8["name"] = img_name |
| x, y = torch.where(window_8 == window_8.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_8["x"] = x |
| max_counts_8["y"] = y |
| |
| if window_14.max().item() > max_counts_14["max"]: |
| max_counts_14["max"] = window_14.max().item() |
| max_counts_14["name"] = img_name |
| x, y = torch.where(window_14 == window_14.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_14["x"] = x |
| max_counts_14["y"] = y |
| |
| if window_16.max().item() > max_counts_16["max"]: |
| max_counts_16["max"] = window_16.max().item() |
| max_counts_16["name"] = img_name |
| x, y = torch.where(window_16 == window_16.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_16["x"] = x |
| max_counts_16["y"] = y |
| |
| if window_28.max().item() > max_counts_28["max"]: |
| max_counts_28["max"] = window_28.max().item() |
| max_counts_28["name"] = img_name |
| x, y = torch.where(window_28 == window_28.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_28["x"] = x |
| max_counts_28["y"] = y |
| |
| if window_32.max().item() > max_counts_32["max"]: |
| max_counts_32["max"] = window_32.max().item() |
| max_counts_32["name"] = img_name |
| x, y = torch.where(window_32 == window_32.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_32["x"] = x |
| max_counts_32["y"] = y |
| |
| if window_56.max().item() > max_counts_56["max"]: |
| max_counts_56["max"] = window_56.max().item() |
| max_counts_56["name"] = img_name |
| x, y = torch.where(window_56 == window_56.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_56["x"] = x |
| max_counts_56["y"] = y |
| |
| if window_64.max().item() > max_counts_64["max"]: |
| max_counts_64["max"] = window_64.max().item() |
| max_counts_64["name"] = img_name |
| x, y = torch.where(window_64 == window_64.max()) |
| x, y = x[0].item(), y[0].item() |
| max_counts_64["x"] = x |
| max_counts_64["y"] = y |
|
|
| window_4 = window_4.view(-1).cpu().numpy().astype(int) |
| window_7 = window_7.view(-1).cpu().numpy().astype(int) |
| window_8 = window_8.view(-1).cpu().numpy().astype(int) |
| window_14 = window_14.view(-1).cpu().numpy().astype(int) |
| window_16 = window_16.view(-1).cpu().numpy().astype(int) |
| window_28 = window_28.view(-1).cpu().numpy().astype(int) |
| window_32 = window_32.view(-1).cpu().numpy().astype(int) |
| window_56 = window_56.view(-1).cpu().numpy().astype(int) |
| window_64 = window_64.view(-1).cpu().numpy().astype(int) |
| |
|
|
| uniques_, counts_ = np.unique(window_4, return_counts=True) |
| counts_4 = _update_dict(counts_4, uniques_, counts_) |
|
|
| uniques_, counts_ = np.unique(window_7, return_counts=True) |
| counts_7 = _update_dict(counts_7, uniques_, counts_) |
|
|
| uniques_, counts_ = np.unique(window_8, return_counts=True) |
| counts_8 = _update_dict(counts_8, uniques_, counts_) |
|
|
| uniques_, counts_ = np.unique(window_14, return_counts=True) |
| counts_14 = _update_dict(counts_14, uniques_, counts_) |
|
|
| uniques_, counts_ = np.unique(window_16, return_counts=True) |
| counts_16 = _update_dict(counts_16, uniques_, counts_) |
|
|
| uniques_, counts_ = np.unique(window_28, return_counts=True) |
| counts_28 = _update_dict(counts_28, uniques_, counts_) |
|
|
| uniques_, counts_ = np.unique(window_32, return_counts=True) |
| counts_32 = _update_dict(counts_32, uniques_, counts_) |
|
|
| uniques_, counts_ = np.unique(window_56, return_counts=True) |
| counts_56 = _update_dict(counts_56, uniques_, counts_) |
|
|
| uniques_, counts_ = np.unique(window_64, return_counts=True) |
| counts_64 = _update_dict(counts_64, uniques_, counts_) |
|
|
| counts = { |
| 1: counts_1, |
| 4: counts_4, |
| 7: counts_7, |
| 8: counts_8, |
| 14: counts_14, |
| 16: counts_16, |
| 28: counts_28, |
| 32: counts_32, |
| 56: counts_56, |
| 64: counts_64 |
| } |
|
|
| max_counts = { |
| 4: max_counts_4, |
| 7: max_counts_7, |
| 8: max_counts_8, |
| 14: max_counts_14, |
| 16: max_counts_16, |
| 28: max_counts_28, |
| 32: max_counts_32, |
| 56: max_counts_56, |
| 64: max_counts_64 |
| } |
|
|
| with open(os.path.join(counts_dir, f"{dataset_name}.json"), "w") as f: |
| json.dump(counts, f) |
| |
| with open(os.path.join(counts_dir, f"{dataset_name}_max.json"), "w") as f: |
| json.dump(max_counts, f) |
|
|
|
|
| def parse_args(): |
| parser = ArgumentParser(description="Get local counts of the dataset") |
| parser.add_argument( |
| "--dataset", |
| type=str, |
| choices=["nwpu", "ucf_qnrf", "shanghaitech_a", "shanghaitech_b"], |
| required=True, |
| help="The dataset to use." |
| ) |
| parser.add_argument( |
| "--device", |
| type=str, |
| default="cuda", |
| help="The device to use." |
| ) |
| args = parser.parse_args() |
| return args |
|
|
|
|
| if __name__ == "__main__": |
| args = parse_args() |
| args.dataset = datasets.standardize_dataset_name(args.dataset) |
| args.device = torch.device(args.device) |
| _get_counts(args.dataset, args.device) |
|
|
