| | 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) |
| |
|
