P2PNet_Soy.ipynb

{"cells":[{"cell_type":"code","execution_count":null,"metadata":{"id":"KwBYMnSdD4hX","colab":{"base_uri":"https://localhost:8080/"},"outputId":"1491974c-9cd6-435f-b2a3-fec78c5eb69f","executionInfo":{"status":"ok","timestamp":1667883249745,"user_tz":-540,"elapsed":183307,"user":{"displayName":"jiangsan zhao","userId":"15833650277394623841"}}},"outputs":[{"output_type":"stream","name":"stdout","text":["Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n","Collecting hdf5storage\n","  Downloading hdf5storage-0.1.18-py2.py3-none-any.whl (53 kB)\n","\u001b[K     |████████████████████████████████| 53 kB 399 kB/s \n","\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from hdf5storage) (1.21.6)\n","Requirement already satisfied: h5py>=2.1 in /usr/local/lib/python3.7/dist-packages (from hdf5storage) (3.1.0)\n","Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py>=2.1->hdf5storage) (1.5.2)\n","Installing collected packages: hdf5storage\n","Successfully installed hdf5storage-0.1.18\n","Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n","Looking in links: https://download.pytorch.org/whl/torch_stable.html\n","Collecting torch==1.8.0+cu111\n","  Downloading https://download.pytorch.org/whl/cu111/torch-1.8.0%2Bcu111-cp37-cp37m-linux_x86_64.whl (1982.2 MB)\n","\u001b[K     |█████████████▌                  | 834.1 MB 1.2 MB/s eta 0:15:22tcmalloc: large alloc 1147494400 bytes == 0x3935a000 @  0x7f269407d615 0x58ead6 0x4f355e 0x4d222f 0x51041f 0x5b4ee6 0x58ff2e 0x510325 0x5b4ee6 0x58ff2e 0x50d482 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4bac0a 0x538a76 0x590ae5 0x510280 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6 0x58ff2e 0x50c4fc 0x58fd37 0x50ca37 0x5b4ee6 0x58ff2e\n","\u001b[K     |█████████████████               | 1055.7 MB 135.1 MB/s eta 0:00:07tcmalloc: large alloc 1434370048 bytes == 0x7d9b0000 @  0x7f269407d615 0x58ead6 0x4f355e 0x4d222f 0x51041f 0x5b4ee6 0x58ff2e 0x510325 0x5b4ee6 0x58ff2e 0x50d482 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4bac0a 0x538a76 0x590ae5 0x510280 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6 0x58ff2e 0x50c4fc 0x58fd37 0x50ca37 0x5b4ee6 0x58ff2e\n","\u001b[K     |█████████████████████▋          | 1336.2 MB 75.2 MB/s eta 0:00:09tcmalloc: large alloc 1792966656 bytes == 0x27e2000 @  0x7f269407d615 0x58ead6 0x4f355e 0x4d222f 0x51041f 0x5b4ee6 0x58ff2e 0x510325 0x5b4ee6 0x58ff2e 0x50d482 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4bac0a 0x538a76 0x590ae5 0x510280 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6 0x58ff2e 0x50c4fc 0x58fd37 0x50ca37 0x5b4ee6 0x58ff2e\n","\u001b[K     |███████████████████████████▎    | 1691.1 MB 1.2 MB/s eta 0:04:06tcmalloc: large alloc 2241208320 bytes == 0x6d5ca000 @  0x7f269407d615 0x58ead6 0x4f355e 0x4d222f 0x51041f 0x5b4ee6 0x58ff2e 0x510325 0x5b4ee6 0x58ff2e 0x50d482 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4d00fb 0x50cb8d 0x4bac0a 0x538a76 0x590ae5 0x510280 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6 0x58ff2e 0x50c4fc 0x58fd37 0x50ca37 0x5b4ee6 0x58ff2e\n","\u001b[K     |████████████████████████████████| 1982.2 MB 1.5 MB/s eta 0:00:01tcmalloc: large alloc 1982251008 bytes == 0xf2f2c000 @  0x7f269407c1e7 0x4b2590 0x4b261c 0x58ead6 0x4f355e 0x4d222f 0x51041f 0x5b4ee6 0x58ff2e 0x50ca37 0x5b4ee6 0x58ff2e 0x50ca37 0x5b4ee6 0x58ff2e 0x50ca37 0x5b4ee6 0x58ff2e 0x50ca37 0x5b4ee6 0x58ff2e 0x50ca37 0x58fd37 0x50ca37 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6\n","tcmalloc: large alloc 2477817856 bytes == 0x1dd64e000 @  0x7f269407d615 0x58ead6 0x4f355e 0x4d222f 0x51041f 0x5b4ee6 0x58ff2e 0x50ca37 0x5b4ee6 0x58ff2e 0x50ca37 0x5b4ee6 0x58ff2e 0x50ca37 0x5b4ee6 0x58ff2e 0x50ca37 0x5b4ee6 0x58ff2e 0x50ca37 0x58fd37 0x50ca37 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6 0x58ff2e 0x50d482 0x5b4ee6 0x4bad99 0x4d3249\n","\u001b[K     |████████████████████████████████| 1982.2 MB 5.3 kB/s \n","\u001b[?25hCollecting torchvision==0.9.0+cu111\n","  Downloading https://download.pytorch.org/whl/cu111/torchvision-0.9.0%2Bcu111-cp37-cp37m-linux_x86_64.whl (17.6 MB)\n","\u001b[K     |████████████████████████████████| 17.6 MB 26.5 MB/s \n","\u001b[?25hCollecting torchaudio==0.8.0\n","  Downloading torchaudio-0.8.0-cp37-cp37m-manylinux1_x86_64.whl (1.9 MB)\n","\u001b[K     |████████████████████████████████| 1.9 MB 4.3 MB/s \n","\u001b[?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.8.0+cu111) (4.1.1)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch==1.8.0+cu111) (1.21.6)\n","Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.7/dist-packages (from torchvision==0.9.0+cu111) (7.1.2)\n","Installing collected packages: torch, torchvision, torchaudio\n","  Attempting uninstall: torch\n","    Found existing installation: torch 1.12.1+cu113\n","    Uninstalling torch-1.12.1+cu113:\n","      Successfully uninstalled torch-1.12.1+cu113\n","  Attempting uninstall: torchvision\n","    Found existing installation: torchvision 0.13.1+cu113\n","    Uninstalling torchvision-0.13.1+cu113:\n","      Successfully uninstalled torchvision-0.13.1+cu113\n","  Attempting uninstall: torchaudio\n","    Found existing installation: torchaudio 0.12.1+cu113\n","    Uninstalling torchaudio-0.12.1+cu113:\n","      Successfully uninstalled torchaudio-0.12.1+cu113\n","\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n","torchtext 0.13.1 requires torch==1.12.1, but you have torch 1.8.0+cu111 which is incompatible.\u001b[0m\n","Successfully installed torch-1.8.0+cu111 torchaudio-0.8.0 torchvision-0.9.0+cu111\n","Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n","Looking in links: https://download.pytorch.org/whl/nightly/cu101/torch-1.7.0.dev20200626%2Bcu101-cp36-cp36m-linux_x86_64.whl\n","Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (1.8.0+cu111)\n","Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch) (4.1.1)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch) (1.21.6)\n","Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n","Collecting tensorboardX\n","  Downloading tensorboardX-2.5.1-py2.py3-none-any.whl (125 kB)\n","\u001b[K     |████████████████████████████████| 125 kB 4.1 MB/s \n","\u001b[?25hRequirement already satisfied: protobuf<=3.20.1,>=3.8.0 in /usr/local/lib/python3.7/dist-packages (from tensorboardX) (3.17.3)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from tensorboardX) (1.21.6)\n","Requirement already satisfied: six>=1.9 in /usr/local/lib/python3.7/dist-packages (from protobuf<=3.20.1,>=3.8.0->tensorboardX) (1.15.0)\n","Installing collected packages: tensorboardX\n","Successfully installed tensorboardX-2.5.1\n"]}],"source":["!pip install hdf5storage\n","!pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html\n","!pip install --pre torch -f  https://download.pytorch.org/whl/nightly/cu101/torch-1.7.0.dev20200626%2Bcu101-cp36-cp36m-linux_x86_64.whl\n","!pip install tensorboardX"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"1laN8KjnD5Lv","colab":{"base_uri":"https://localhost:8080/"},"outputId":"9bb92917-8637-4b92-ea2d-389d1416787e","executionInfo":{"status":"ok","timestamp":1667883419123,"user_tz":-540,"elapsed":169385,"user":{"displayName":"jiangsan zhao","userId":"15833650277394623841"}}},"outputs":[{"output_type":"stream","name":"stdout","text":["Mounted at /content/drive\n"]}],"source":["from google.colab import drive\n","drive.mount('/content/drive')\n","import os\n","os.chdir(\"/content/drive/My Drive\")"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"d2dc1-u6KMTS"},"outputs":[],"source":["## \n","import os\n","import random\n","from scipy import spatial\n","import networkx as nx\n","\n","import torch\n","import numpy as np\n","from torch.utils.data import Dataset\n","from PIL import Image\n","import cv2\n","import glob\n","import scipy.io as io\n","from matplotlib import pyplot as plt\n","plt.switch_backend('agg')\n","\n","class SHHA(Dataset):\n","    def __init__(self, data_root, transform=None, train=False, patch=False, flip=False):\n","        self.root_path = data_root\n","        self.train_lists = os.path.join(self.root_path, \"soybean_seed_counting_a.txt\")\n","        self.eval_list = os.path.join(self.root_path, \"soybean_seed_counting_b.txt\")\n","        # \n","        if train:\n","            self.img_list_file = [name.split(',') for name in open(self.train_lists).read().splitlines()]\n","        else:\n","            self.img_list_file = [name.split(',') for name in open(self.eval_list).read().splitlines()]\n","\n","        self.img_list = self.img_list_file\n","        \n","        # \n","        self.nSamples = len(self.img_list)\n","        \n","        self.transform = transform\n","        self.train = train\n","        self.patch = patch\n","        self.flip = flip\n","\n","    def __len__(self):\n","        return self.nSamples\n","\n","    def __getitem__(self, index):\n","        assert index <= len(self), 'index range error'\n","\n","        img_path = self.img_list[index][0]\n","        gt_path = self.img_list[index][1]\n","        # \n","        img, point = load_data((img_path, gt_path), self.train)\n","        #\n","        if self.transform is not None:\n","            img = self.transform(img)\n","\n","        if self.train:\n","            # data augmentation -> random scale\n","            scale_range = [0.5, 1.4]\n","            min_size = min(img.shape[1:])\n","            scale = random.uniform(*scale_range)\n","            # scale the image and points\n","            if scale * min_size > 224:\n","                img = torch.nn.functional.upsample_bilinear(img.unsqueeze(0), scale_factor=scale).squeeze(0)\n","                point *= scale\n","        # random crop augumentaiton\n","        if self.train and self.patch:\n","            img, point = random_crop(img, point)\n","            for i, _ in enumerate(point):\n","                point[i] = torch.Tensor(point[i])\n","        # random flipping\n","        if random.random() > 0.1 and self.train and self.flip: # never flip\n","            # random flip\n","            img = torch.Tensor(img[:, :, :, ::-1].copy())\n","            for i, _ in enumerate(point):\n","                point[i][:, 0] = 224 - point[i][:, 0]\n","        # random change brightness\n","        if random.random() > 0.3 and self.train: # never flip\n","            #\n","            img = (torch.Tensor(img).clone())*random.uniform(8,12)/10\n","            for i, _ in enumerate(point):\n","                point[i][:, 0] = point[i][:, 0]\n","\n","        if not self.train:\n","            point = [point]\n","\n","        img = torch.Tensor(img)\n","        #  need to adapt your own image names\n","        target = [{} for i in range(len(point))]\n","        for i, _ in enumerate(point):\n","            target[i]['point'] = torch.Tensor(point[i])\n","            image_id_1 = int(img_path.split('/')[-1].split('.')[0].split(\"_\")[1][4:8])\n","            image_id_1 = torch.Tensor([image_id_1]).long()\n","            #\n","            image_id_2 = int(img_path.split('/')[-1].split('.')[0].split(\"_\")[3])\n","            image_id_2 = torch.Tensor([image_id_2]).long()\n","            target[i]['image_id_1'] = image_id_1\n","            target[i]['image_id_2'] = image_id_2\n","            target[i]['labels'] = torch.ones([point[i].shape[0]]).long()\n","\n","        return img, target\n","\n","\n","def load_data(img_gt_path, train):\n","    img_path, gt_path = img_gt_path\n","    # load the images\n","    img = cv2.imread(img_path)\n","    img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n","    # load ground truth points\n","    points = []\n","    #\n","    pts = open(gt_path).read().splitlines()\n","    for pt_0 in pts:\n","        pt = eval(pt_0)        \n","        x = float(pt[0])\n","        y = float(pt[1])\n","        points.append([x, y])\n","    return img, np.array(points)\n","\n","# random crop augumentation\n","def random_crop(img, den, num_patch=10):\n","    half_h = 224\n","    half_w = 224\n","    result_img = np.zeros([num_patch, img.shape[0], half_h, half_w])\n","    result_den = []\n","    # \n","    for i in range(num_patch):\n","        start_h = random.randint(0, img.size(1) - half_h)\n","        start_w = random.randint(0, img.size(2) - half_w)\n","        end_h = start_h + half_h\n","        end_w = start_w + half_w\n","        # \n","        result_img[i] = img[:, start_h:end_h, start_w:end_w]#*random.uniform(5,15)/10\n","        # copy the cropped points\n","        idx = (den[:, 0] >= start_w) & (den[:, 0] <= end_w) & (den[:, 1] >= start_h) & (den[:, 1] <= end_h)\n","        # \n","        record_den = den[idx]\n","        record_den[:, 0] -= start_w\n","        record_den[:, 1] -= start_h\n","\n","        result_den.append(record_den)\n","\n","    return result_img, result_den"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"_ZHJfljJc-bb"},"outputs":[],"source":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"8fkDa6XEKMWE"},"outputs":[],"source":["# \n","import torchvision.transforms as standard_transforms\n","\n","# \n","class DeNormalize(object):\n","    def __init__(self, mean, std):\n","        self.mean = mean\n","        self.std = std\n","\n","    def __call__(self, tensor):\n","        for t, m, s in zip(tensor, self.mean, self.std):\n","            t.mul_(s).add_(m)\n","        return tensor\n","\n","def loading_data(data_root):\n","    # \n","    transform = standard_transforms.Compose([\n","        standard_transforms.ToTensor(), \n","        standard_transforms.Normalize(mean=[0.485, 0.456, 0.406],\n","                                    std=[0.229, 0.224, 0.225]),\n","    ])\n","    # \n","    train_set = SHHA(data_root, train=True, transform=transform, patch=True, flip=True)\n","    # \n","    val_set = SHHA(data_root, train=False, transform=transform)\n","\n","    return train_set, val_set"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"q32CayMIjYwX"},"outputs":[],"source":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"wZY5NX0-KMZO"},"outputs":[],"source":["# \n","import math\n","import os\n","import sys\n","from typing import Iterable\n","import torch\n","import numpy as np\n","import time\n","import torchvision.transforms as standard_transforms\n","import cv2\n","\n","class DeNormalize(object):\n","    def __init__(self, mean, std):\n","        self.mean = mean\n","        self.std = std\n","\n","    def __call__(self, tensor):\n","        for t, m, s in zip(tensor, self.mean, self.std):\n","            t.mul_(s).add_(m)\n","        return tensor\n","#\n","def vis(samples, targets, pred, vis_dir, epoch, predict_cnt, gt_cnt):\n","    '''\n","    samples -> tensor: [batch, 3, H, W]\n","    targets -> list of dict: [{'points':[], 'image_id': str}]\n","    pred -> list: [num_preds, 2]\n","    '''\n","    gts = [t['point'].tolist() for t in targets]\n","\n","    pil_to_tensor = standard_transforms.ToTensor()\n","\n","    restore_transform = standard_transforms.Compose([\n","        DeNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n","        standard_transforms.ToPILImage()\n","    ])\n","    # \n","    for idx in range(samples.shape[0]):\n","        sample = restore_transform(samples[idx])\n","        sample = pil_to_tensor(sample.convert('RGB')).numpy() * 255\n","        sample_gt = sample.transpose([1, 2, 0])[:, :, :].astype(np.uint8).copy()\n","        sample_pred = sample.transpose([1, 2, 0])[:, :, :].astype(np.uint8).copy()\n","\n","        max_len = np.max(sample_gt.shape)\n","\n","        size = 5\n","        # draw gt\n","        for t in gts[idx]:\n","            sample_gt = cv2.circle(sample_gt, (int(t[0]), int(t[1])), size, (0, 255, 0), -1)\n","        # draw predictions\n","        for p in pred[idx]:\n","            sample_pred = cv2.circle(sample_pred, (int(p[0]), int(p[1])), size, (0, 0, 255), -1)\n","\n","        name_1 = targets[idx]['image_id_1']\n","        name_2 = targets[idx]['image_id_2']\n","        #################\n","        fig = plt.figure()\n","        ax1 = fig.add_subplot(1, 2, 1)\n","        ax1.imshow(sample_gt)\n","        ax1.get_xaxis().set_visible(False)\n","        ax1.get_yaxis().set_visible(False)\n","        ax2 = fig.add_subplot(1, 2, 2)\n","        ax2.imshow(sample_pred)\n","        ax2.get_xaxis().set_visible(False)\n","        ax2.get_yaxis().set_visible(False)\n","        fig.suptitle('manual count=%4.2f, inferred count=%4.2f'%(gt_cnt, predict_cnt), fontsize=10)\n","        plt.tight_layout(rect=[0, 0, 0.95, 0.95]) # maize tassels counting\n","        plt.savefig(os.path.join(vis_dir, '{}_{}_id_{}_ind_{}.jpg'.format(epoch, idx, int(name_1), int(name_2))), bbox_inches='tight', dpi = 300)\n","        plt.close()\n","# the training\n","def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,\n","                    data_loader: Iterable, optimizer: torch.optim.Optimizer,\n","                    device: torch.device, epoch: int, max_norm: float = 0):\n","    model.train()\n","    criterion.train()\n","    metric_logger = MetricLogger(delimiter=\"  \")\n","    metric_logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}'))\n","    # iterate all training samples\n","    for samples, targets in data_loader:\n","        #\n","        samples = samples.to(device)\n","        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n","        # forward\n","        outputs = model(samples)\n","        #\n","        # calc the losses\n","        loss_dict = criterion(outputs, targets)\n","        weight_dict = criterion.weight_dict\n","        losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)\n","        # reduce all losses (get the mean values)\n","        loss_dict_reduced = reduce_dict(loss_dict)\n","        loss_dict_reduced_unscaled = {f'{k}_unscaled': v\n","                                      for k, v in loss_dict_reduced.items()}\n","        loss_dict_reduced_scaled = {k: v * weight_dict[k]\n","                                    for k, v in loss_dict_reduced.items() if k in weight_dict}\n","        losses_reduced_scaled = sum(loss_dict_reduced_scaled.values())\n","\n","        loss_value = losses_reduced_scaled.item()\n","\n","        if not math.isfinite(loss_value):\n","            print(\"Loss is {}, stopping training\".format(loss_value))\n","            print(loss_dict_reduced)\n","            sys.exit(1)\n","        # backward\n","        optimizer.zero_grad()\n","        losses.backward()\n","        if max_norm > 0:\n","            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)\n","        optimizer.step()\n","        # update logger\n","        metric_logger.update(loss=loss_value, **loss_dict_reduced_scaled, **loss_dict_reduced_unscaled)\n","        metric_logger.update(lr=optimizer.param_groups[0][\"lr\"])\n","    # gather the stats from all processes\n","    metric_logger.synchronize_between_processes()\n","    print(\"Averaged stats:\", metric_logger)\n","    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}\n","\n","# evaluate the model performance during training\n","@torch.no_grad()\n","def evaluate_crowd_no_overlap(model, data_loader, device, epoch, threshold, vis_dir=None):\n","    model.eval()\n","    metric_logger = MetricLogger(delimiter=\"  \")\n","    metric_logger.add_meter('class_error', SmoothedValue(window_size=1, fmt='{value:.2f}'))\n","    # run inference on all images to calc MAE\n","    maes = []\n","    mses = []\n","    for samples, targets in data_loader:\n","\n","        samples = samples.to(device)\n","\n","        outputs = model(samples)\n","        outputs_scores = torch.nn.functional.softmax(outputs['pred_logits'], -1)[:, :, 1][0]\n","\n","        outputs_points = outputs['pred_points'][0]\n","\n","        gt_cnt = targets[0]['point'].shape[0]\n","        # 0.5 is used by default\n","        threshold = threshold\n","\n","        points = outputs_points[outputs_scores > threshold].detach().cpu().numpy()#.tolist()\n","        # choose to merge closely located points\n","        if points.shape[0]<10000 and points.shape[0] != 0:\n","            # choose the cut off point\n","            cutoff = 500/points.shape[0]\n","            if cutoff<20:\n","                cutoff = 20\n","            components = nx.connected_components(\n","                nx.from_edgelist(\n","                    (i, j) for i, js in enumerate(\n","                        spatial.KDTree(points).query_ball_point(points, cutoff)\n","                    )\n","                    for j in js\n","                )\n","            )\n","\n","            clusters = {j: i for i, js in enumerate(components) for j in js}\n","\n","            # reorganize the points to the order of clusters \n","            points_reo = np.zeros(points.shape)\n","            i = 0\n","            for key in clusters.keys():\n","                points_reo[i,:] = points[key,:]\n","                i+=1\n","            # points_n has the same order as clusters\n","            res = [clusters[key] for key in clusters.keys()]\n","            res_n = np.array(res).reshape(-1,1)\n","\n","            points_n = []\n","            for i in np.unique(res_n):\n","                tmp = points_reo[np.where(res_n[:,0] == i)]\n","                points_n.append( [np.mean(tmp[:,0]), np.mean(tmp[:,1])])\n","        else:\n","            points_n = points.tolist()\n","\n","        predict_cnt = len(points_n)\n","        #save the visualized images\n","        if vis_dir is not None: \n","            vis(samples, targets, [points_n], vis_dir, epoch, predict_cnt, gt_cnt)\n","        # accumulate MAE, MSE\n","        mae = abs(predict_cnt - gt_cnt)\n","        mse = (predict_cnt - gt_cnt) * (predict_cnt - gt_cnt)\n","        maes.append(float(mae))\n","        mses.append(float(mse))\n","    # calc MAE, MSE\n","    mae = np.mean(maes)\n","    mse = np.sqrt(np.mean(mses))\n","\n","    return mae, mse"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"6GqceuB4CGnu"},"outputs":[],"source":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"yB4BETJp2m6W"},"outputs":[],"source":["import argparse\n","import datetime\n","import random\n","import time\n","from pathlib import Path\n","from IPython.display import clear_output \n","\n","import torch\n","from torch.utils.data import DataLoader, DistributedSampler\n","\n","import os\n","from tensorboardX import SummaryWriter\n","import warnings\n","warnings.filterwarnings('ignore')\n"]},{"cell_type":"code","source":["# \n","import os\n","import subprocess\n","import time\n","from collections import defaultdict, deque\n","import datetime\n","import pickle\n","from typing import Optional, List\n","\n","import torch\n","import torch.distributed as dist\n","from torch import Tensor\n","\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torch.autograd import Variable\n","\n","# needed due to empty tensor bug in pytorch and torchvision 0.5\n","import torchvision\n","if float(torchvision.__version__[:3]) < 0.7:\n","    from torchvision.ops import _new_empty_tensor\n","    from torchvision.ops.misc import _output_size\n","\n","\n","class SmoothedValue(object):\n","    \"\"\"Track a series of values and provide access to smoothed values over a\n","    window or the global series average.\n","    \"\"\"\n","    def __init__(self, window_size=20, fmt=None):\n","        if fmt is None:\n","            fmt = \"{median:.4f} ({global_avg:.4f})\"\n","        self.deque = deque(maxlen=window_size)\n","        self.total = 0.0\n","        self.count = 0\n","        self.fmt = fmt\n","\n","    def update(self, value, n=1):\n","        self.deque.append(value)\n","        self.count += n\n","        self.total += value * n\n","\n","    def synchronize_between_processes(self):\n","        \"\"\"\n","        Warning: does not synchronize the deque!\n","        \"\"\"\n","        if not is_dist_avail_and_initialized():\n","            return\n","        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')\n","        dist.barrier()\n","        dist.all_reduce(t)\n","        t = t.tolist()\n","        self.count = int(t[0])\n","        self.total = t[1]\n","\n","    @property\n","    def median(self):\n","        d = torch.tensor(list(self.deque))\n","        return d.median().item()\n","\n","    @property\n","    def avg(self):\n","        d = torch.tensor(list(self.deque), dtype=torch.float32)\n","        return d.mean().item()\n","\n","    @property\n","    def global_avg(self):\n","        return self.total / self.count\n","\n","    @property\n","    def max(self):\n","        return max(self.deque)\n","\n","    @property\n","    def value(self):\n","        return self.deque[-1]\n","\n","    def __str__(self):\n","        return self.fmt.format(\n","            median=self.median,\n","            avg=self.avg,\n","            global_avg=self.global_avg,\n","            max=self.max,\n","            value=self.value)\n","#\n","def all_gather(data):\n","    \"\"\"\n","    Run all_gather on arbitrary picklable data (not necessarily tensors)\n","    Args:\n","        data: any picklable object\n","    Returns:\n","        list[data]: list of data gathered from each rank\n","    \"\"\"\n","    world_size = get_world_size()\n","    if world_size == 1:\n","        return [data]\n","\n","    # serialized to a Tensor\n","    buffer = pickle.dumps(data)\n","    storage = torch.ByteStorage.from_buffer(buffer)\n","    tensor = torch.ByteTensor(storage).to(\"cuda\")\n","\n","    # obtain Tensor size of each rank\n","    local_size = torch.tensor([tensor.numel()], device=\"cuda\")\n","    size_list = [torch.tensor([0], device=\"cuda\") for _ in range(world_size)]\n","    dist.all_gather(size_list, local_size)\n","    size_list = [int(size.item()) for size in size_list]\n","    max_size = max(size_list)\n","\n","    # receiving Tensor from all ranks\n","    # we pad the tensor because torch all_gather does not support\n","    # gathering tensors of different shapes\n","    tensor_list = []\n","    for _ in size_list:\n","        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device=\"cuda\"))\n","    if local_size != max_size:\n","        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device=\"cuda\")\n","        tensor = torch.cat((tensor, padding), dim=0)\n","    dist.all_gather(tensor_list, tensor)\n","\n","    data_list = []\n","    for size, tensor in zip(size_list, tensor_list):\n","        buffer = tensor.cpu().numpy().tobytes()[:size]\n","        data_list.append(pickle.loads(buffer))\n","\n","    return data_list\n","#\n","def reduce_dict(input_dict, average=True):\n","    \"\"\"\n","    Args:\n","        input_dict (dict): all the values will be reduced\n","        average (bool): whether to do average or sum\n","    Reduce the values in the dictionary from all processes so that all processes\n","    have the averaged results. Returns a dict with the same fields as\n","    input_dict, after reduction.\n","    \"\"\"\n","    world_size = get_world_size()\n","    if world_size < 2:\n","        return input_dict\n","    with torch.no_grad():\n","        names = []\n","        values = []\n","        # sort the keys so that they are consistent across processes\n","        for k in sorted(input_dict.keys()):\n","            names.append(k)\n","            values.append(input_dict[k])\n","        values = torch.stack(values, dim=0)\n","        dist.all_reduce(values)\n","        if average:\n","            values /= world_size\n","        reduced_dict = {k: v for k, v in zip(names, values)}\n","    return reduced_dict\n","#\n","class MetricLogger(object):\n","    def __init__(self, delimiter=\"\\t\"):\n","        self.meters = defaultdict(SmoothedValue)\n","        self.delimiter = delimiter\n","\n","    def update(self, **kwargs):\n","        for k, v in kwargs.items():\n","            if isinstance(v, torch.Tensor):\n","                v = v.item()\n","            assert isinstance(v, (float, int))\n","            self.meters[k].update(v)\n","\n","    def __getattr__(self, attr):\n","        if attr in self.meters:\n","            return self.meters[attr]\n","        if attr in self.__dict__:\n","            return self.__dict__[attr]\n","        raise AttributeError(\"'{}' object has no attribute '{}'\".format(\n","            type(self).__name__, attr))\n","\n","    def __str__(self):\n","        loss_str = []\n","        for name, meter in self.meters.items():\n","            loss_str.append(\n","                \"{}: {}\".format(name, str(meter))\n","            )\n","        return self.delimiter.join(loss_str)\n","\n","    def synchronize_between_processes(self):\n","        for meter in self.meters.values():\n","            meter.synchronize_between_processes()\n","\n","    def add_meter(self, name, meter):\n","        self.meters[name] = meter\n","\n","    def log_every(self, iterable, print_freq, header=None):\n","        i = 0\n","        if not header:\n","            header = ''\n","        start_time = time.time()\n","        end = time.time()\n","        iter_time = SmoothedValue(fmt='{avg:.4f}')\n","        data_time = SmoothedValue(fmt='{avg:.4f}')\n","        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'\n","        if torch.cuda.is_available():\n","            log_msg = self.delimiter.join([\n","                header,\n","                '[{0' + space_fmt + '}/{1}]',\n","                'eta: {eta}',\n","                '{meters}',\n","                'time: {time}',\n","                'data: {data}',\n","                'max mem: {memory:.0f}'\n","            ])\n","        else:\n","            log_msg = self.delimiter.join([\n","                header,\n","                '[{0' + space_fmt + '}/{1}]',\n","                'eta: {eta}',\n","                '{meters}',\n","                'time: {time}',\n","                'data: {data}'\n","            ])\n","        MB = 1024.0 * 1024.0\n","        for obj in iterable:\n","            data_time.update(time.time() - end)\n","            yield obj\n","            iter_time.update(time.time() - end)\n","            if i % print_freq == 0 or i == len(iterable) - 1:\n","                eta_seconds = iter_time.global_avg * (len(iterable) - i)\n","                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))\n","                if torch.cuda.is_available():\n","                    print(log_msg.format(\n","                        i, len(iterable), eta=eta_string,\n","                        meters=str(self),\n","                        time=str(iter_time), data=str(data_time),\n","                        memory=torch.cuda.max_memory_allocated() / MB))\n","                else:\n","                    print(log_msg.format(\n","                        i, len(iterable), eta=eta_string,\n","                        meters=str(self),\n","                        time=str(iter_time), data=str(data_time)))\n","            i += 1\n","            end = time.time()\n","        total_time = time.time() - start_time\n","        total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n","        print('{} Total time: {} ({:.4f} s / it)'.format(\n","            header, total_time_str, total_time / len(iterable)))\n","#\n","def get_sha():\n","    cwd = os.path.dirname(os.path.abspath(__file__))\n","\n","    def _run(command):\n","        return subprocess.check_output(command, cwd=cwd).decode('ascii').strip()\n","    sha = 'N/A'\n","    diff = \"clean\"\n","    branch = 'N/A'\n","    try:\n","        sha = _run(['git', 'rev-parse', 'HEAD'])\n","        subprocess.check_output(['git', 'diff'], cwd=cwd)\n","        diff = _run(['git', 'diff-index', 'HEAD'])\n","        diff = \"has uncommited changes\" if diff else \"clean\"\n","        branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])\n","    except Exception:\n","        pass\n","    message = f\"sha: {sha}, status: {diff}, branch: {branch}\"\n","    return message\n","\n","\n","def collate_fn(batch):\n","    batch = list(zip(*batch))\n","    batch[0] = nested_tensor_from_tensor_list(batch[0])\n","    return tuple(batch)\n","\n","def collate_fn_crowd(batch):\n","    # re-organize the batch\n","    batch_new = []\n","    for b in batch:\n","        imgs, points = b\n","        if imgs.ndim == 3:\n","            imgs = imgs.unsqueeze(0)\n","        for i in range(len(imgs)):\n","            batch_new.append((imgs[i, :, :, :], points[i]))\n","    batch = batch_new\n","    batch = list(zip(*batch))\n","    batch[0] = nested_tensor_from_tensor_list(batch[0])\n","    return tuple(batch)\n","\n","\n","def _max_by_axis(the_list):\n","    # type: (List[List[int]]) -> List[int]\n","    maxes = the_list[0]\n","    for sublist in the_list[1:]:\n","        for index, item in enumerate(sublist):\n","            maxes[index] = max(maxes[index], item)\n","    return maxes\n","\n","def _max_by_axis_pad(the_list):\n","    # type: (List[List[int]]) -> List[int]\n","    maxes = the_list[0]\n","    for sublist in the_list[1:]:\n","        for index, item in enumerate(sublist):\n","            maxes[index] = max(maxes[index], item)\n","\n","    block = 128\n","\n","    for i in range(2):\n","        maxes[i+1] = ((maxes[i+1] - 1) // block + 1) * block\n","    return maxes\n","#\n","def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):\n","    # TODO make this more general\n","    if tensor_list[0].ndim == 3:\n","\n","        # TODO make it support different-sized images\n","        max_size = _max_by_axis_pad([list(img.shape) for img in tensor_list])\n","        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))\n","        batch_shape = [len(tensor_list)] + max_size\n","        b, c, h, w = batch_shape\n","        dtype = tensor_list[0].dtype\n","        device = tensor_list[0].device\n","        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)\n","        for img, pad_img in zip(tensor_list, tensor):\n","            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)\n","    else:\n","        raise ValueError('not supported')\n","    return tensor\n","#\n","class NestedTensor(object):\n","    def __init__(self, tensors, mask: Optional[Tensor]):\n","        self.tensors = tensors\n","        self.mask = mask\n","\n","    def to(self, device):\n","        # type: (Device) -> NestedTensor # noqa\n","        cast_tensor = self.tensors.to(device)\n","        mask = self.mask\n","        if mask is not None:\n","            assert mask is not None\n","            cast_mask = mask.to(device)\n","        else:\n","            cast_mask = None\n","        return NestedTensor(cast_tensor, cast_mask)\n","\n","    def decompose(self):\n","        return self.tensors, self.mask\n","\n","    def __repr__(self):\n","        return str(self.tensors)\n","#\n","def setup_for_distributed(is_master):\n","    \"\"\"\n","    This function disables printing when not in master process\n","    \"\"\"\n","    import builtins as __builtin__\n","    builtin_print = __builtin__.print\n","\n","    def print(*args, **kwargs):\n","        force = kwargs.pop('force', False)\n","        if is_master or force:\n","            builtin_print(*args, **kwargs)\n","\n","    __builtin__.print = print\n","#\n","def is_dist_avail_and_initialized():\n","    if not dist.is_available():\n","        return False\n","    if not dist.is_initialized():\n","        return False\n","    return True\n","#\n","def get_world_size():\n","    if not is_dist_avail_and_initialized():\n","        return 1\n","    return dist.get_world_size()\n","#\n","def get_rank():\n","    if not is_dist_avail_and_initialized():\n","        return 0\n","    return dist.get_rank()\n","#\n","def is_main_process():\n","    return get_rank() == 0\n","#\n","def save_on_master(*args, **kwargs):\n","    if is_main_process():\n","        torch.save(*args, **kwargs)\n","#\n","def init_distributed_mode(args):\n","    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:\n","        args.rank = int(os.environ[\"RANK\"])\n","        args.world_size = int(os.environ['WORLD_SIZE'])\n","        args.gpu = int(os.environ['LOCAL_RANK'])\n","    elif 'SLURM_PROCID' in os.environ:\n","        args.rank = int(os.environ['SLURM_PROCID'])\n","        args.gpu = args.rank % torch.cuda.device_count()\n","    else:\n","        print('Not using distributed mode')\n","        args.distributed = False\n","        return\n","\n","    args.distributed = True\n","\n","    torch.cuda.set_device(args.gpu)\n","    args.dist_backend = 'nccl'\n","    print('| distributed init (rank {}): {}'.format(\n","        args.rank, args.dist_url), flush=True)\n","    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,\n","                                         world_size=args.world_size, rank=args.rank)\n","    torch.distributed.barrier()\n","    setup_for_distributed(args.rank == 0)\n","#\n","@torch.no_grad()\n","def accuracy(output, target, topk=(1,)):\n","    \"\"\"Computes the precision@k for the specified values of k\"\"\"\n","    if target.numel() == 0:\n","        return [torch.zeros([], device=output.device)]\n","    maxk = max(topk)\n","    batch_size = target.size(0)\n","\n","    _, pred = output.topk(maxk, 1, True, True)\n","    pred = pred.t()\n","    correct = pred.eq(target.view(1, -1).expand_as(pred))\n","\n","    res = []\n","    for k in topk:\n","        correct_k = correct[:k].view(-1).float().sum(0)\n","        res.append(correct_k.mul_(100.0 / batch_size))\n","    return res\n","#\n","def interpolate(input, size=None, scale_factor=None, mode=\"nearest\", align_corners=None):\n","    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor\n","    \"\"\"\n","    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.\n","    This will eventually be supported natively by PyTorch, and this\n","    class can go away.\n","    \"\"\"\n","    if float(torchvision.__version__[:3]) < 0.7:\n","        if input.numel() > 0:\n","            return torch.nn.functional.interpolate(\n","                input, size, scale_factor, mode, align_corners\n","            )\n","\n","        output_shape = _output_size(2, input, size, scale_factor)\n","        output_shape = list(input.shape[:-2]) + list(output_shape)\n","        return _new_empty_tensor(input, output_shape)\n","    else:\n","        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)\n","#\n","class FocalLoss(nn.Module):\n","    r\"\"\"\n","        This criterion is a implemenation of Focal Loss, which is proposed in\n","        Focal Loss for Dense Object Detection.\n","\n","            Loss(x, class) = - \\alpha (1-softmax(x)[class])^gamma \\log(softmax(x)[class])\n","\n","        The losses are averaged across observations for each minibatch.\n","\n","        Args:\n","            alpha(1D Tensor, Variable) : the scalar factor for this criterion\n","            gamma(float, double) : gamma > 0; reduces the relative loss for well-classiﬁed examples (p > .5),\n","                                   putting more focus on hard, misclassiﬁed examples\n","            size_average(bool): By default, the losses are averaged over observations for each minibatch.\n","                                However, if the field size_average is set to False, the losses are\n","                                instead summed for each minibatch.\n","\n","\n","    \"\"\"\n","    def __init__(self, class_num, alpha=None, gamma=2, size_average=True):\n","        super(FocalLoss, self).__init__()\n","        if alpha is None:\n","            self.alpha = Variable(torch.ones(class_num, 1))\n","        else:\n","            if isinstance(alpha, Variable):\n","                self.alpha = alpha\n","            else:\n","                self.alpha = Variable(alpha)\n","        self.gamma = gamma\n","        self.class_num = class_num\n","        self.size_average = size_average\n","\n","    def forward(self, inputs, targets):\n","        N = inputs.size(0)\n","        C = inputs.size(1)\n","        P = F.softmax(inputs)\n","\n","        class_mask = inputs.data.new(N, C).fill_(0)\n","        class_mask = Variable(class_mask)\n","        ids = targets.view(-1, 1)\n","        class_mask.scatter_(1, ids.data, 1.)\n","\n","        if inputs.is_cuda and not self.alpha.is_cuda:\n","            self.alpha = self.alpha.cuda()\n","        alpha = self.alpha[ids.data.view(-1)]\n","\n","        probs = (P*class_mask).sum(1).view(-1,1)\n","\n","        log_p = probs.log()\n","        batch_loss = -alpha*(torch.pow((1-probs), self.gamma))*log_p\n","\n","        if self.size_average:\n","            loss = batch_loss.mean()\n","        else:\n","            loss = batch_loss.sum()\n","        return loss"],"metadata":{"id":"mPv72RxO2u-n"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["## \n","import torch\n","from scipy.optimize import linear_sum_assignment\n","from torch import nn\n","\n","class HungarianMatcher_Crowd(nn.Module):\n","    \"\"\"This class computes an assignment between the targets and the predictions of the network\n","\n","    For efficiency reasons, the targets don't include the no_object. Because of this, in general,\n","    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,\n","    while the others are un-matched (and thus treated as non-objects).\n","    \"\"\"\n","\n","    def __init__(self, cost_class: float = 1, cost_point: float = 1):\n","        \"\"\"Creates the matcher\n","\n","        Params:\n","            cost_class: This is the relative weight of the foreground object\n","            cost_point: This is the relative weight of the L1 error of the points coordinates in the matching cost\n","        \"\"\"\n","        super().__init__()\n","        self.cost_class = cost_class\n","        self.cost_point = cost_point\n","        assert cost_class != 0 or cost_point != 0, \"all costs cant be 0\"\n","\n","    @torch.no_grad()\n","    def forward(self, outputs, targets):\n","        \"\"\" Performs the matching\n","\n","        Params:\n","            outputs: This is a dict that contains at least these entries:\n","                 \"pred_logits\": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits\n","                 \"points\": Tensor of dim [batch_size, num_queries, 2] with the predicted point coordinates\n","\n","            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:\n","                 \"labels\": Tensor of dim [num_target_points] (where num_target_points is the number of ground-truth\n","                           objects in the target) containing the class labels\n","                 \"points\": Tensor of dim [num_target_points, 2] containing the target point coordinates\n","\n","        Returns:\n","            A list of size batch_size, containing tuples of (index_i, index_j) where:\n","                - index_i is the indices of the selected predictions (in order)\n","                - index_j is the indices of the corresponding selected targets (in order)\n","            For each batch element, it holds:\n","                len(index_i) = len(index_j) = min(num_queries, num_target_points)\n","        \"\"\"\n","        bs, num_queries = outputs[\"pred_logits\"].shape[:2]\n","\n","        # We flatten to compute the cost matrices in a batch\n","        out_prob = outputs[\"pred_logits\"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]\n","        out_points = outputs[\"pred_points\"].flatten(0, 1)  # [batch_size * num_queries, 2]\n","\n","        # Also concat the target labels and points\n","        # tgt_ids = torch.cat([v[\"labels\"] for v in targets])\n","        tgt_ids = torch.cat([v[\"labels\"] for v in targets])\n","        tgt_points = torch.cat([v[\"point\"] for v in targets])\n","\n","        # Compute the classification cost. Contrary to the loss, we don't use the NLL,\n","        # but approximate it in 1 - proba[target class].\n","        # The 1 is a constant that doesn't change the matching, it can be ommitted.\n","        cost_class = -out_prob[:, tgt_ids]\n","\n","        # Compute the L2 cost between point\n","        cost_point = torch.cdist(out_points, tgt_points, p=2)\n","\n","        # Compute the giou cost between point\n","\n","        # Final cost matrix\n","        C = self.cost_point * cost_point + self.cost_class * cost_class\n","        C = C.view(bs, num_queries, -1).cpu()\n","\n","        sizes = [len(v[\"point\"]) for v in targets]\n","        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]\n","        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]\n","#\n","def build_matcher_crowd(args):\n","    return HungarianMatcher_Crowd(cost_class=args.set_cost_class, cost_point=args.set_cost_point)"],"metadata":{"id":"DkJ2o7YK_1hO"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["### The feature extraction part is mainly adapted from https://github.com/zhaoyuzhi/PyTorch-Pyramid-Feature-Attention-Network-for-Saliency-Detection"],"metadata":{"id":"pBf2_5jxO-D7"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"L669uVsdCvFy"},"outputs":[],"source":["## \n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","#\n","class SpatialAttention(nn.Module):\n","    def __init__(self, in_channels, kernel_size=9):\n","        super(SpatialAttention, self).__init__()\n","\n","        self.kernel_size = kernel_size\n","        self.in_channels = in_channels\n","        pad = (self.kernel_size-1)//2  # Padding on one side for stride 1\n","\n","        self.grp1_conv1k = nn.Conv2d(self.in_channels, self.in_channels//2, (1, self.kernel_size), padding=(0, pad))\n","        self.grp1_bn1 = nn.BatchNorm2d(self.in_channels//2)\n","        self.grp1_convk1 = nn.Conv2d(self.in_channels//2, 1, (self.kernel_size, 1), padding=(pad, 0))\n","        self.grp1_bn2 = nn.BatchNorm2d(1)\n","\n","        self.grp2_convk1 = nn.Conv2d(self.in_channels, self.in_channels//2, (self.kernel_size, 1), padding=(pad, 0))\n","        self.grp2_bn1 = nn.BatchNorm2d(self.in_channels//2)\n","        self.grp2_conv1k = nn.Conv2d(self.in_channels//2, 1, (1, self.kernel_size), padding=(0, pad))\n","        self.grp2_bn2 = nn.BatchNorm2d(1)\n","\n","    def forward(self, input_):\n","        # Generate Group 1 Features\n","        grp1_feats = self.grp1_conv1k(input_)\n","        grp1_feats = F.relu(self.grp1_bn1(grp1_feats))\n","        grp1_feats = self.grp1_convk1(grp1_feats)\n","        grp1_feats = F.relu(self.grp1_bn2(grp1_feats))\n","\n","        # Generate Group 2 features\n","        grp2_feats = self.grp2_convk1(input_)\n","        grp2_feats = F.relu(self.grp2_bn1(grp2_feats))\n","        grp2_feats = self.grp2_conv1k(grp2_feats)\n","        grp2_feats = F.relu(self.grp2_bn2(grp2_feats))\n","\n","        added_feats = torch.sigmoid(torch.add(grp1_feats, grp2_feats))\n","        added_feats = added_feats.expand_as(input_).clone()\n","\n","        return added_feats\n","#\n","class ChannelwiseAttention(nn.Module):\n","    def __init__(self, in_channels):\n","        super(ChannelwiseAttention, self).__init__()\n","\n","        self.in_channels = in_channels\n","\n","        self.linear_1 = nn.Linear(self.in_channels, self.in_channels//4)\n","        self.linear_2 = nn.Linear(self.in_channels//4, self.in_channels)\n","\n","    def forward(self, input_):\n","        n_b, n_c, h, w = input_.size()\n","\n","        feats = F.adaptive_avg_pool2d(input_, (1, 1)).view((n_b, n_c))\n","        feats = F.relu(self.linear_1(feats))\n","        feats = torch.sigmoid(self.linear_2(feats))\n","        \n","        # Activity regularizer\n","        ca_act_reg = torch.mean(feats)\n","\n","        feats = feats.view((n_b, n_c, 1, 1))\n","        feats = feats.expand_as(input_).clone()\n","\n","        return feats, ca_act_reg"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"VmPunNWUGSOe"},"outputs":[],"source":["#### \n","import numpy as np\n","import cv2\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import torchvision.models as models\n","\n","vgg_conv1_2 = vgg_conv2_2 = vgg_conv3_3 = vgg_conv4_3 = vgg_conv5_3 = None\n","\n","def conv_1_2_hook(module, input, output):\n","    global vgg_conv1_2\n","    vgg_conv1_2 = output\n","    return None\n","\n","def conv_2_2_hook(module, input, output):\n","    global vgg_conv2_2\n","    vgg_conv2_2 = output\n","    return None\n","\n","def conv_3_3_hook(module, input, output):\n","    global vgg_conv3_3\n","    vgg_conv3_3 = output\n","    return None\n","\n","def conv_4_3_hook(module, input, output):\n","    global vgg_conv4_3\n","    vgg_conv4_3 = output\n","    return None\n","\n","def conv_5_3_hook(module, input, output):\n","    global vgg_conv5_3\n","    vgg_conv5_3 = output\n","    return None\n","\n","##\n","class CPFE_hl(nn.Module):\n","    def __init__(self, feature_layer=None, out_channels=8):\n","        super(CPFE_hl, self).__init__()\n","\n","        self.dil_rates = [3, 5, 7]\n","\n","        # Determine number of in_channels from VGG-16 feature layer\n","        if feature_layer == 'conv5_3':\n","            self.in_channels = 512\n","        elif feature_layer == 'conv4_3':\n","            self.in_channels = 512\n","        elif feature_layer == 'conv3_3':\n","            self.in_channels = 256\n","        elif feature_layer == 'conv2_3':\n","            self.in_channels = 128\n","        elif feature_layer == 'conv1_3':\n","            self.in_channels = 64\n","\n","        # Define layers\n","        self.conv_1_1 = nn.Conv2d(in_channels=self.in_channels, out_channels=out_channels, kernel_size=1, bias=False)\n","        self.conv_dil_3 = nn.Conv2d(in_channels=self.in_channels, out_channels=out_channels, kernel_size=3,\n","                                    stride=1, dilation=self.dil_rates[0], padding=self.dil_rates[0], bias=False)\n","        self.conv_dil_5 = nn.Conv2d(in_channels=self.in_channels, out_channels=out_channels, kernel_size=3,\n","                                    stride=1, dilation=self.dil_rates[1], padding=self.dil_rates[1], bias=False)\n","        self.conv_dil_7 = nn.Conv2d(in_channels=self.in_channels, out_channels=out_channels, kernel_size=3,\n","                                    stride=1, dilation=self.dil_rates[2], padding=self.dil_rates[2], bias=False)\n","\n","        self.bn = nn.BatchNorm2d(out_channels*4)\n","\n","    def forward(self, input_):\n","        # Extract features\n","        conv_1_1_feats = self.conv_1_1(input_)\n","        conv_dil_3_feats = self.conv_dil_3(input_)\n","        conv_dil_5_feats = self.conv_dil_5(input_)\n","        conv_dil_7_feats = self.conv_dil_7(input_)\n","\n","        # Aggregate features\n","        concat_feats = torch.cat((conv_1_1_feats, conv_dil_3_feats, conv_dil_5_feats, conv_dil_7_feats), dim=1)\n","        bn_feats = F.relu(self.bn(concat_feats))\n","\n","        return bn_feats\n","##\n","class SODModel(nn.Module):\n","    def __init__(self):\n","        super(SODModel, self).__init__()\n","\n","        # \n","        self.vgg16 = models.vgg16(pretrained=True).features\n","\n","        # Extract and register intermediate features of VGG-16_bn\n","        self.vgg16[3].register_forward_hook(conv_1_2_hook)\n","        self.vgg16[8].register_forward_hook(conv_2_2_hook)\n","        self.vgg16[15].register_forward_hook(conv_3_3_hook)\n","        self.vgg16[22].register_forward_hook(conv_4_3_hook)\n","        self.vgg16[29].register_forward_hook(conv_5_3_hook)\n","\n","        # Initialize layers for high level (hl) feature (conv3_3, conv4_3, conv5_3) processing\n","        self.cpfe_conv3_3 = CPFE_hl(feature_layer='conv3_3')\n","        self.cpfe_conv4_3 = CPFE_hl(feature_layer='conv4_3')\n","        self.cpfe_conv5_3 = CPFE_hl(feature_layer='conv5_3')\n","        #\n","        self.cpfe_conv1_3 = CPFE_hl(feature_layer='conv1_3')\n","        self.cpfe_conv2_3 = CPFE_hl(feature_layer='conv2_3')\n","        # 11,03,2022, remove channel attention\n","        self.cha_att = ChannelwiseAttention(in_channels=96)  # in_channels = 3 x (8 x 4)\n","\n","        self.hl_conv1 = nn.Conv2d(96, 8, (3, 3), padding=1)\n","        self.hl_bn1 = nn.BatchNorm2d(8)\n","\n","        # \n","        self.ll_conv_1 = nn.Conv2d(64, 8, (3, 3), padding=1)\n","        self.ll_bn_1 = nn.BatchNorm2d(8)\n","        self.ll_conv_2 = nn.Conv2d(128, 8, (3, 3), padding=1)\n","        self.ll_bn_2 = nn.BatchNorm2d(8)\n","        self.ll_conv_3 = nn.Conv2d(64, 8, (3, 3), padding=1) \n","        self.ll_bn_3 = nn.BatchNorm2d(8)\n","\n","        self.spa_att = SpatialAttention(in_channels=8)\n","\n","        # \n","        self.ff_conv_1 = nn.Conv2d(16, 3, (3, 3), padding=1)\n","        self.ff_bn_1 = nn.BatchNorm2d(3)\n","    def forward(self, input_):\n","        global vgg_conv1_2, vgg_conv2_2, vgg_conv3_3, vgg_conv4_3, vgg_conv5_3\n","\n","        # Pass input_ through vgg16 to generate intermediate features\n","        self.vgg16(input_)\n","        # Process high level features\n","        conv3_cpfe_feats = self.cpfe_conv3_3(vgg_conv3_3)\n","        conv4_cpfe_feats = self.cpfe_conv4_3(vgg_conv4_3)\n","        conv5_cpfe_feats = self.cpfe_conv5_3(vgg_conv5_3)\n","\n","        conv4_cpfe_feats = F.interpolate(conv4_cpfe_feats, scale_factor=2, mode='bilinear', align_corners=True) # reduce spatial dimension by 2\n","        conv5_cpfe_feats = F.interpolate(conv5_cpfe_feats, scale_factor=4, mode='bilinear', align_corners=True)\n","\n","        conv_345_feats = torch.cat((conv3_cpfe_feats, conv4_cpfe_feats, conv5_cpfe_feats), dim=1)\n","        \n","        # channel attention on high level features\n","        conv_345_ca, ca_act_reg = self.cha_att(conv_345_feats)\n","        conv_345_feats = torch.mul(conv_345_feats, conv_345_ca)\n","\n","        conv_345_feats = self.hl_conv1(conv_345_feats)\n","        conv_345_feats = F.relu(self.hl_bn1(conv_345_feats))\n","        ##\n","        # Process low level features\n","        conv0_feats = input_ # the original input image\n","        conv1_cpfe_feats = self.cpfe_conv1_3(vgg_conv1_2)\n","        conv2_cpfe_feats = self.cpfe_conv2_3(vgg_conv2_2)\n","\n","        conv0_feats = F.interpolate(conv0_feats, scale_factor=0.25, mode='bilinear', align_corners=True)\n","        conv1_cpfe_feats = F.interpolate(conv1_cpfe_feats, scale_factor=0.25, mode='bilinear', align_corners=True)\n","        conv2_cpfe_feats = F.interpolate(conv2_cpfe_feats, scale_factor=0.5, mode='bilinear', align_corners=True)\n","\n","        #\n","        conv_12_feats = torch.cat((conv1_cpfe_feats, conv2_cpfe_feats), dim=1)\n","        conv_12_feats = self.ll_conv_3(conv_12_feats)\n","        conv_12_feats = F.relu(self.ll_bn_3(conv_12_feats))\n","        # spatial attention on low level features\n","        conv_12_sa = self.spa_att(conv_12_feats)\n","        conv_12_feats = torch.mul(conv_12_feats, conv_12_sa)\n","\n","        # fuse the low and high level features\n","        fused_feats = torch.cat((conv_12_feats, conv_345_feats), dim=1)\n","        #\n","        fused_final = self.ff_conv_1(fused_feats)\n","        fused_final = F.relu(self.ff_bn_1(fused_final))\n","        # add the fused low and high level features to the original image\n","        fused_final_out = torch.add(fused_final, conv0_feats)\n","        fused_final_out = torch.sigmoid(fused_final_out)\n","\n","        return fused_final_out\n"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"jNz8iOaxdj6d"},"outputs":[],"source":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Y3M_OTlkVPO5"},"outputs":[],"source":["# build model p2pNet.py\n","import torch\n","import torch.nn.functional as F\n","from torch import nn\n","import numpy as np\n","import time\n","\n","# the network frmawork of the regression branch\n","class RegressionModel(nn.Module):\n","    def __init__(self, num_features_in, num_anchor_points=4, feature_size=32):\n","        super(RegressionModel, self).__init__()\n","\n","        self.conv1 = nn.Conv2d(num_features_in, feature_size, kernel_size=3, padding=1)\n","        self.act1 = nn.ReLU()\n","\n","        self.conv2 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)\n","        self.act2 = nn.ReLU()\n","\n","        self.conv3 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)\n","        self.act3 = nn.ReLU()\n","\n","        self.conv4 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)\n","        self.act4 = nn.ReLU()\n","\n","        self.output = nn.Conv2d(feature_size, num_anchor_points * 2, kernel_size=3, padding=1) # one point has two coordinates \n","    # sub-branch forward\n","    def forward(self, x):\n","        out = self.conv1(x)\n","        out = self.act1(out)\n","\n","        out = self.conv2(out)\n","        out = self.act2(out)\n","\n","        out = self.output(out)\n","\n","        out = out.permute(0, 2, 3, 1)\n","\n","        return out.contiguous().view(out.shape[0], -1, 2)\n","\n","# the network frmawork of the classification branch\n","class ClassificationModel(nn.Module):\n","    def __init__(self, num_features_in, num_anchor_points=4, num_classes=80, prior=0.01, feature_size=32):\n","        super(ClassificationModel, self).__init__()\n","\n","        self.num_classes = num_classes\n","        self.num_anchor_points = num_anchor_points\n","\n","        self.conv1 = nn.Conv2d(num_features_in, feature_size, kernel_size=3, padding=1)\n","        self.act1 = nn.ReLU()\n","\n","        self.conv2 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)\n","        self.act2 = nn.ReLU()\n","\n","        self.conv3 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)\n","        self.act3 = nn.ReLU()\n","\n","        self.conv4 = nn.Conv2d(feature_size, feature_size, kernel_size=3, padding=1)\n","        self.act4 = nn.ReLU()\n","\n","        self.output = nn.Conv2d(feature_size, num_anchor_points * num_classes, kernel_size=3, padding=1) # one classes, only positives \n","        self.output_act = nn.Sigmoid()\n","    # sub-branch forward\n","    def forward(self, x):\n","        out = self.conv1(x)\n","        out = self.act1(out)\n","\n","        out = self.conv2(out)\n","        out = self.act2(out)\n","\n","        out = self.output(out)\n","\n","        out1 = out.permute(0, 2, 3, 1)\n","\n","        batch_size, width, height, _ = out1.shape\n","\n","        out2 = out1.view(batch_size, width, height, self.num_anchor_points, self.num_classes)\n","\n","        return out2.contiguous().view(x.shape[0], -1, self.num_classes)\n","\n","# generate the reference points in grid layout\n","def generate_anchor_points(stride=8, row=3, line=3):\n","    row_step = stride / row\n","    line_step = stride / line\n","\n","    shift_x = (np.arange(1, line + 1) - 0.5) * line_step - stride / 2\n","    shift_y = (np.arange(1, row + 1) - 0.5) * row_step - stride / 2\n","\n","    shift_x, shift_y = np.meshgrid(shift_x, shift_y)\n","\n","    anchor_points = np.vstack((\n","        shift_x.ravel(), shift_y.ravel()\n","    )).transpose()\n","\n","    return anchor_points\n","# shift the meta-anchor to get an acnhor points\n","def shift(shape, stride, anchor_points):\n","    shift_x = (np.arange(0, shape[1]) + 0.5)* stride\n","    shift_y = (np.arange(0, shape[0]) + 0.5)* stride\n","\n","    shift_x, shift_y = np.meshgrid(shift_x, shift_y)\n","\n","    shifts = np.vstack((\n","        shift_x.ravel(), shift_y.ravel()\n","    )).transpose()\n","\n","    A = anchor_points.shape[0]\n","    K = shifts.shape[0]\n","    all_anchor_points = (anchor_points.reshape((1, A, 2)) + shifts.reshape((1, K, 2)).transpose((1, 0, 2)))\n","    all_anchor_points = all_anchor_points.reshape((K * A, 2))\n","\n","    return all_anchor_points\n","\n","# \n","class AnchorPoints(nn.Module):\n","    def __init__(self, pyramid_levels=None, strides=None, row=3, line=3):\n","        super(AnchorPoints, self).__init__()\n","\n","        if pyramid_levels is None:\n","            self.pyramid_levels = [3, 4, 5, 6, 7]\n","        else:\n","            self.pyramid_levels = pyramid_levels\n","\n","        if strides is None:\n","            self.strides = [2 ** x for x in self.pyramid_levels]\n","\n","        self.row = row\n","        self.line = line\n","\n","    def forward(self, image):\n","        image_shape = image.shape[2:]\n","        image_shape = np.array(image_shape)\n","        image_shapes = [(image_shape + 2 ** x - 1) // (2 ** x) for x in self.pyramid_levels] # calcualtes the output size of the model (image of 128*128 to feature map of 16*16)\n","\n","        all_anchor_points = np.zeros((0, 2)).astype(np.float32)\n","        # get reference points for each level\n","        for idx, p in enumerate(self.pyramid_levels):\n","            anchor_points = generate_anchor_points(2**p, row=self.row, line=self.line)\n","            shifted_anchor_points = shift(image_shapes[idx], self.strides[idx], anchor_points)\n","            all_anchor_points = np.append(all_anchor_points, shifted_anchor_points, axis=0)\n","\n","        all_anchor_points = np.expand_dims(all_anchor_points, axis=0)\n","        # send reference points to device\n","        if torch.cuda.is_available():\n","            return torch.from_numpy(all_anchor_points.astype(np.float32)).cuda()\n","        else:\n","            return torch.from_numpy(all_anchor_points.astype(np.float32))\n","##\n","# the defenition of the P2PNet model\n","class P2PNet(nn.Module):\n","    def __init__(self, row=2, line=2):\n","        super().__init__()\n","        self.num_classes = 2\n","        # the number of all anchor points\n","        num_anchor_points = row * line\n","\n","        self.regression = RegressionModel(num_features_in=3, num_anchor_points=num_anchor_points)\n","        self.classification = ClassificationModel(num_features_in=3, \\\n","                                            num_classes=self.num_classes, \\\n","                                            num_anchor_points=num_anchor_points)\n","\n","        self.anchor_points = AnchorPoints(pyramid_levels=[2,], row=row, line=line) # remember to change pyramid level when you change feature input\n","\n","        self.fpn = SODModel()\n","\n","    def forward(self, samples: NestedTensor): #\n","        # \n","        features_fpn = self.fpn(samples) # output = bach_size, channel, Height, Weight\n","        #\n","        batch_size = features_fpn.size()[0]\n","        # run the regression and classification branch\n","        regression = self.regression(features_fpn) * 100 # 8x\n","        classification = self.classification(features_fpn)\n","        #\n","        anchor_points = self.anchor_points(samples).repeat(batch_size, 1, 1)\n","        output_coord = regression + anchor_points\n","        output_class = classification\n","        out = {'pred_logits': output_class, 'pred_points': output_coord}\n","        #\n","        return out\n","\n","class SetCriterion_Crowd(nn.Module):\n","\n","    def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses):\n","        \"\"\" Create the criterion.\n","        Parameters:\n","            num_classes: number of object categories, omitting the special no-object category\n","            matcher: module able to compute a matching between targets and proposals\n","            weight_dict: dict containing as key the names of the losses and as values their relative weight.\n","            eos_coef: relative classification weight applied to the no-object category\n","            losses: list of all the losses to be applied. See get_loss for list of available losses.\n","        \"\"\"\n","        super().__init__()\n","        self.num_classes = num_classes\n","        self.matcher = matcher\n","        self.weight_dict = weight_dict\n","        self.eos_coef = eos_coef\n","        self.losses = losses\n","        empty_weight = torch.ones(self.num_classes + 1)\n","        empty_weight[0] = self.eos_coef\n","        self.register_buffer('empty_weight', empty_weight)\n","\n","    def loss_labels(self, outputs, targets, indices, num_points):\n","        \"\"\"Classification loss (NLL)\n","        targets dicts must contain the key \"labels\" containing a tensor of dim [nb_target_boxes]\n","        \"\"\"\n","        assert 'pred_logits' in outputs\n","        src_logits = outputs['pred_logits']\n","\n","        idx = self._get_src_permutation_idx(indices)\n","        target_classes_o = torch.cat([t[\"labels\"][J] for t, (_, J) in zip(targets, indices)])\n","        target_classes = torch.full(src_logits.shape[:2], 0,\n","                                    dtype=torch.int64, device=src_logits.device)\n","        target_classes[idx] = target_classes_o\n","\n","        loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)\n","        losses = {'loss_ce': loss_ce}\n","\n","        return losses\n","\n","    def loss_points(self, outputs, targets, indices, num_points):\n","\n","        assert 'pred_points' in outputs\n","        idx = self._get_src_permutation_idx(indices)\n","        src_points = outputs['pred_points'][idx]\n","        target_points = torch.cat([t['point'][i] for t, (_, i) in zip(targets, indices)], dim=0)\n","        #print(\"target_points {}\".format(target_points))\n","        loss_bbox = F.mse_loss(src_points, target_points, reduction='none')\n","\n","        losses = {}\n","        losses['loss_points'] = loss_bbox.sum() / num_points\n","\n","        return losses\n","\n","    def _get_src_permutation_idx(self, indices):\n","        # permute predictions following indices\n","        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])\n","        src_idx = torch.cat([src for (src, _) in indices])\n","        return batch_idx, src_idx\n","\n","    def _get_tgt_permutation_idx(self, indices):\n","        # permute targets following indices\n","        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])\n","        tgt_idx = torch.cat([tgt for (_, tgt) in indices])\n","        return batch_idx, tgt_idx\n","\n","    def get_loss(self, loss, outputs, targets, indices, num_points, **kwargs):\n","        loss_map = {\n","            'labels': self.loss_labels,\n","            'points': self.loss_points,\n","        }\n","        #print(\"loss_map {}\".format(loss_map))\n","        assert loss in loss_map, f'do you really want to compute {loss} loss?'\n","        return loss_map[loss](outputs, targets, indices, num_points, **kwargs)\n","\n","    def forward(self, outputs, targets):\n","        \"\"\" This performs the loss computation.\n","        Parameters:\n","             outputs: dict of tensors, see the output specification of the model for the format\n","             targets: list of dicts, such that len(targets) == batch_size.\n","                      The expected keys in each dict depends on the losses applied, see each loss' doc\n","        \"\"\"\n","        output1 = {'pred_logits': outputs['pred_logits'], 'pred_points': outputs['pred_points']}\n","\n","        indices1 = self.matcher(output1, targets)\n","\n","        num_points = sum(len(t[\"labels\"]) for t in targets)\n","\n","        num_points = torch.as_tensor([num_points], dtype=torch.float, device=next(iter(output1.values())).device)\n","\n","        num_boxes = torch.clamp(num_points / get_world_size(), min=1).item()\n","\n","        losses = {}\n","        for loss in self.losses:\n","            losses.update(self.get_loss(loss, output1, targets, indices1, num_boxes)) #\n","\n","        return losses\n","\n","# create the P2PNet model\n","def build(args, training):\n","    # treats persons as a single class\n","    num_classes = 1\n","\n","    model = P2PNet(args.row, args.line)\n","    if not training: \n","        return model\n","\n","    weight_dict = {'loss_ce': 1, 'loss_points': args.point_loss_coef}\n","    losses = ['labels', 'points']     #['labels', 'points']\n","    matcher = build_matcher_crowd(args)\n","    criterion = SetCriterion_Crowd(num_classes, \\\n","                                matcher=matcher, weight_dict=weight_dict, \\\n","                                eos_coef=args.eos_coef, losses=losses)\n","    return model, criterion\n","##\n"]},{"cell_type":"markdown","source":["### List all the parameters here"],"metadata":{"id":"O3Gd0Z3JSDE8"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"zllbyl5BOlQj"},"outputs":[],"source":["\n","def get_arguments():\n","    \"\"\"Parse all the arguments provided from the CLI.\n","    Returns:\n","      A list of parsed arguments.\n","    \"\"\"\n","    parser = argparse.ArgumentParser(description=\"Object Counting Framework\")\n","    # constant\n","    parser.add_argument('--lr', default=1e-3, type=float)\n","    parser.add_argument('--lr_fpn', default=1e-5, type=float)\n","    parser.add_argument('--batch_size', default=1, type=int)\n","    parser.add_argument('--weight_decay', default=1e-4, type=float)\n","    parser.add_argument('--epochs', default=3500, type=int)\n","    parser.add_argument('--lr_drop', default=100, type=int)\n","    parser.add_argument('--clip_max_norm', default=0.1, type=float,\n","                        help='gradient clipping max norm')\n","\n","    # Model parameters\n","    parser.add_argument('--frozen_weights', type=str, default=None,\n","                        help=\"Path to the pretrained model. If set, only the mask head will be trained\")\n","    #\n","    parser.add_argument('--set_cost_class', default=1, type=float,\n","                        help=\"Class coefficient in the matching cost\")\n","\n","    parser.add_argument('--set_cost_point', default=0.99, type=float,\n","                        help=\"L1 point coefficient in the matching cost\")\n","\n","    # * Loss coefficients\n","    parser.add_argument('--point_loss_coef', default=0.02, type=float) # default = 0.0002 # 0.5\n","    parser.add_argument('--eos_coef', default=0.02, type=float, # 0.05\n","                        help=\"Relative classification weight of the no-object class\") # default = 0.5\n","    \n","    # a threshold during evaluation for counting and visualization\n","    parser.add_argument('--threshold', default=0.5, type=float,\n","                        help=\"threshold in evalluation: evaluate_crowd_no_overlap\")\n","    parser.add_argument('--row', default=2, type=int,\n","                        help=\"row number of anchor points\")\n","    parser.add_argument('--line', default=2, type=int,\n","                        help=\"line number of anchor points\")\n","\n","    # dataset parameters\n","    parser.add_argument('--dataset_file', default='SHHA')\n","    parser.add_argument('--data_root', default='/content/drive/My Drive/P2PNet-Soy/Soybean_seed_counting/',\n","                        help='path where the dataset is')\n","    \n","    parser.add_argument('--output_dir', default='/content/drive/My Drive/P2PNet-Soy/log_P2PNet_Soy',\n","                        help='path where to save, empty for no saving')\n","    parser.add_argument('--checkpoints_dir', default='/content/drive/My Drive/P2PNet-Soy/ckpt_P2PNet_Soy_01',\n","                        help='path where to save checkpoints, empty for no saving') \n","    parser.add_argument('--tensorboard_dir', default='/content/drive/My Drive/P2PNet-Soy/runs_P2PNet_Soy',\n","                        help='path where to save, empty for no saving')\n","\n","    parser.add_argument('--seed', default=42, type=int)\n","    parser.add_argument('--resume', default='', help='resume from checkpoint')\n","    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',\n","                        help='start epoch')\n","    parser.add_argument('--eval', action='store_true')\n","    parser.add_argument('--num_workers', default=1, type=int)\n","    parser.add_argument('--eval_freq', default=3, type=int,\n","                        help='frequency of evaluation, default setting is evaluating in every 5 epoch')\n","    parser.add_argument('--gpu_id', default=0, type=int, help='the gpu used for training')\n","    #\n","    opt = parser.parse_known_args()[0]\n","    return opt"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Gl6zkwnxkqJs"},"outputs":[],"source":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"MzU3pfq1Js4r","colab":{"base_uri":"https://localhost:8080/","height":84,"referenced_widgets":["80ec627ed5f3417ebecbbcfe767338aa","ec2555b45230458896dbeb40f581d433","db3be6617bf34876be88384c1a6e9597","fa532929d8774047abdc95b5af3296a9","4be575664ce449ddb5cbef4377147e38","ceed4556f6474d2bbaeba192fa0cb1e5","690f8ed4dcaa4f95b415366319ab5944","44a8014a98b3499692332bf9838f151a","ce00405a13e945aa89616c7947053ae3","675bf8a98fbd4208b1f5537c8fe7ecd4","4c31f05ca00541a2b018c6a8bfc393e6"]},"outputId":"a3821481-41c4-4ea3-8f3a-22dc9629c3ba"},"outputs":[{"output_type":"stream","name":"stderr","text":["Downloading: \"https://download.pytorch.org/models/vgg16-397923af.pth\" to /root/.cache/torch/hub/checkpoints/vgg16-397923af.pth\n"]},{"output_type":"display_data","data":{"text/plain":["  0%|          | 0.00/528M [00:00<?, ?B/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"80ec627ed5f3417ebecbbcfe767338aa"}},"metadata":{}},{"output_type":"stream","name":"stdout","text":["number of params: 15137927\n"]}],"source":["args = get_arguments()\n","#\n","#put the model path here if you have trained any or comment it out\n","args.resume = \"/content/drive/My Drive/P2PNet-Soy/ckpt_P2PNet_Soy/best_mae.pth\" \n","# the directory to save the evaluations during training\n","args.vis_dir = \"/content/drive/My Drive/P2PNet-Soy/vis_P2PNet_Soy\"\n","if not os.path.exists(args.vis_dir):\n","    os.makedirs(args.vis_dir)\n","##\n","os.environ[\"CUDA_VISIBLE_DEVICES\"] = '{}'.format(args.gpu_id)\n","# create the logging file\n","if not os.path.exists(args.output_dir):\n","    os.makedirs(args.output_dir)\n","run_log_name = os.path.join(args.output_dir, 'run_log.txt')\n","with open(run_log_name, \"w\") as log_file:\n","    log_file.write('Eval Log %s\\n' % time.strftime(\"%c\"))\n","#\n","with open(run_log_name, \"a\") as log_file:\n","    log_file.write(\"{}\".format(args))\n","device = torch.device('cuda')\n","# fix the seed for reproducibility\n","seed = args.seed + get_rank()\n","seed = args.seed\n","torch.manual_seed(seed)\n","np.random.seed(seed)\n","random.seed(seed)\n","# get the P2PNet model\n","model, criterion = build(args, training=True)\n","# move to GPU\n","model.to(device)\n","criterion.to(device)\n","\n","model_without_ddp = model\n","\n","n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","print('number of params:', n_parameters)\n","# use different optimation params for different parts of the model\n","param_dicts = [\n","    {\"params\": [p for n, p in model_without_ddp.named_parameters() if \"fpn\" not in n and p.requires_grad]},\n","    {\n","        \"params\": [p for n, p in model_without_ddp.named_parameters() if \"fpn\" in n and p.requires_grad],\n","        \"lr\": args.lr_fpn,\n","    },\n","]\n","# Adam is used by default\n","optimizer = torch.optim.Adam(param_dicts, lr=args.lr)\n","lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)\n","# create the training and valiation set\n","train_set, val_set = loading_data(args.data_root)\n","# create the sampler used during training\n","sampler_train = torch.utils.data.RandomSampler(train_set)\n","sampler_val = torch.utils.data.SequentialSampler(val_set)\n","\n","batch_sampler_train = torch.utils.data.BatchSampler(\n","    sampler_train, args.batch_size, drop_last=True)\n","# the dataloader for training\n","data_loader_train = DataLoader(train_set, batch_sampler=batch_sampler_train,\n","                                collate_fn=collate_fn_crowd, num_workers=args.num_workers)\n","\n","data_loader_val = DataLoader(val_set, 1, sampler=sampler_val,\n","                                drop_last=False, collate_fn=collate_fn_crowd, num_workers=args.num_workers)\n","\n","if args.frozen_weights is not None:\n","    checkpoint = torch.load(args.frozen_weights, map_location='cpu')\n","    model_without_ddp.detr.load_state_dict(checkpoint['model'])\n","# resume the weights and training state if exists\n","if args.resume:\n","    checkpoint = torch.load(args.resume, map_location='cpu')\n","    model_without_ddp.load_state_dict(checkpoint['model'])\n","    args.start_epoch = checkpoint['epoch']\n","    new_start = 1\n","    if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:\n","        optimizer.load_state_dict(checkpoint['optimizer'])\n","        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n","else:\n","    new_start = 0\n","##\n","start_time = time.time()\n","# save the performance during the training\n","mae = []\n","mse = []\n","epoch_save = []\n","# the logger writer\n","writer = SummaryWriter(args.tensorboard_dir)\n","# save latest weights every epoch\n","if not os.path.exists(args.checkpoints_dir):\n","    os.makedirs(args.checkpoints_dir)\n","#\n","step = 0\n","# training starts here\n","for epoch in range(args.start_epoch, args.epochs):\n","    # always run evaluation first (to check what model has been loaded !!!)\n","    if (epoch +2) % args.eval_freq == 0 or new_start: # and epoch != 0\n","        # change the status right after the first iteration\n","        new_start = 0\n","        #\n","        t1 = time.time()\n","        result = evaluate_crowd_no_overlap(model, data_loader_val, device, epoch, args.threshold, args.vis_dir)\n","        t2 = time.time()\n","        print(\"evaluation time {}\".format(t2-t1))\n","        mae.append(result[0])\n","        mse.append(result[1])\n","        epoch_save.append(epoch)\n","        #\n","        epoch_save_m = np.array(epoch_save)[mae == np.min(mae)][0]\n","        # print the evaluation results\n","        print('=======================================test=======================================')\n","        print(\"mae:\", result[0], \"mse:\", result[1], \"time:\", t2 - t1, \"best mae:\", np.min(mae), \"at epoch: {}\".format(epoch_save_m) )\n","        with open(run_log_name, \"a\") as log_file:\n","            log_file.write(\"mae:{}, mse:{}, time:{}, best mae:{}\".format(result[0], \n","                            result[1], t2 - t1, np.min(mae)))\n","        print('=======================================test=======================================')\n","        # recored the evaluation results\n","        if writer is not None:\n","            with open(run_log_name, \"a\") as log_file:\n","                log_file.write(\"metric/mae@{}: {}\".format(step, result[0]))\n","                log_file.write(\"metric/mse@{}: {}\".format(step, result[1]))\n","            writer.add_scalar('metric/mae', result[0], step)\n","            writer.add_scalar('metric/mse', result[1], step)\n","            step += 1\n","\n","        # save the best model since begining\n","        if abs(np.min(mae) - result[0]) < 0.01:\n","            checkpoint_best_path = os.path.join(args.checkpoints_dir, 'best_mae.pth')\n","            torch.save({\n","                'model': model_without_ddp.state_dict(),\n","                'epoch': epoch,\n","            }, checkpoint_best_path)\n","    ###\n","    t1 = time.time()\n","    stat = train_one_epoch(model, criterion, data_loader_train, optimizer, device, epoch, args.clip_max_norm)\n","\n","    # record the training states after every epoch\n","    if writer is not None:\n","        with open(run_log_name, \"a\") as log_file:\n","            log_file.write(\"loss/loss@{}: {}\".format(epoch, stat['loss']))\n","            log_file.write(\"loss/loss_ce@{}: {}\".format(epoch, stat['loss_ce']))\n","            \n","        writer.add_scalar('loss/loss', stat['loss'], epoch)\n","        writer.add_scalar('loss/loss_ce', stat['loss_ce'], epoch)\n","\n","    t2 = time.time()\n","    print('[ep %d][lr %.7f][%.2fs]' % \\\n","            (epoch, optimizer.param_groups[0]['lr'], t2 - t1))\n","    with open(run_log_name, \"a\") as log_file:\n","        log_file.write('[ep %d][lr %.7f][%.2fs]' % (epoch, optimizer.param_groups[0]['lr'], t2 - t1))\n","    # change lr according to the scheduler\n","    lr_scheduler.step()\n","    #\n","    # save latest weights every epoch\n","    checkpoint_latest_path = os.path.join(args.checkpoints_dir, 'latest.pth')\n","    torch.save({\n","        'model': model_without_ddp.state_dict(),\n","        'epoch': epoch,\n","    }, checkpoint_latest_path)\n","    ## clear the cell output regulary\n","    if epoch % 150 == 0 and epoch != 0:\n","        clear_output()\n","# total time for training\n","total_time = time.time() - start_time\n","total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n","print('Training time {}'.format(total_time_str))\n"]},{"cell_type":"code","source":[],"metadata":{"id":"a4L6DvrcUdzh"},"execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"AydizQosJs-B"},"outputs":[],"source":[]},{"cell_type":"code","source":[],"metadata":{"id":"9xF_Yj6l29BG"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["##Inference"],"metadata":{"id":"4Pk9IBmU29nk"}},{"cell_type":"code","source":["# create the P2PNet model\n","def build_eval(args):\n","    model = P2PNet(args.row, args.line)\n","    return model"],"metadata":{"id":"WyPH9FTt29EB"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["def get_arguments():\n","    \"\"\"Parse all the arguments provided from the CLI.\n","    Returns:\n","      A list of parsed arguments.\n","    \"\"\"\n","    parser = argparse.ArgumentParser(description=\"Object Counting Framework\")\n","    # constant\n","    # a threshold during evaluation for counting and visualization\n","    parser.add_argument('--threshold', default=0.5, type=float,\n","                        help=\"threshold in evalluation: evaluate_crowd_no_overlap\")\n","    parser.add_argument('--row', default=2, type=int,\n","                        help=\"row number of anchor points\")\n","    parser.add_argument('--line', default=2, type=int,\n","                        help=\"line number of anchor points\")\n","    parser.add_argument('--data_root', default='/content/drive/My Drive/soypod_crop_counting/',\n","                        help='path where the dataset is')\n","    parser.add_argument('--seed', default=42, type=int)\n","    parser.add_argument('--resume', default='', help='resume from checkpoint')\n","    parser.add_argument('--num_workers', default=1, type=int)\n","    parser.add_argument('--gpu_id', default=0, type=int, help='the gpu used for training')\n","    #\n","    opt = parser.parse_known_args()[0] #if known else parser.parse_args()\n","    return opt"],"metadata":{"id":"qKBGhFOn29Gr"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["args = get_arguments()\n","# specify the directories to model weights and prediction output\n","args.resume = \"/content/drive/My Drive/P2PNet-Soy/ckpt_P2PNet_Soy/best_mae.pth\"\n","args.vis_dir = \"/content/drive/My Drive/P2PNet-Soy/vis_P2PNet_Soy_out\""],"metadata":{"id":"CQY3kYKE8IxF"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["####\n","if not os.path.exists(args.vis_dir):\n","    os.makedirs(args.vis_dir)\n","##\n","os.environ[\"CUDA_VISIBLE_DEVICES\"] = '{}'.format(args.gpu_id)\n","\n","\n","device = torch.device('cuda')\n","# fix the seed for reproducibility\n","seed = args.seed + get_rank()\n","random.seed(seed)\n","# original model\n","model = build_eval(args, training = False)\n","# move to GPU\n","model.to(device)"],"metadata":{"id":"On6TPlsN8g4u"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["###On images of cropped individual plant"],"metadata":{"id":"t2wq7UoA8lpA"}},{"cell_type":"code","source":["\n","# threshold for evaluation\n","threshold = 0.5 #args.threshold\n","# to apply post processing or not: if Filter = 10000, the filter is applied\n","\n","#\n","# load trained model\n","if args.resume is not None:\n","    checkpoint = torch.load(args.resume, map_location='cpu')\n","    model.load_state_dict(checkpoint['model'])\n","#\n","model.eval()\n","#\n","step = 0\n","epoch = 0\n","###\n","# create the pre-processing transform\n","transform = standard_transforms.Compose([\n","    standard_transforms.ToTensor(), \n","    standard_transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n","])\n","\n","# set your image path here\n","for img_type in [\"b\", \"a\"]: #\n","    print(\"current image folder is {}\".format(img_type))\n","    img_list_path = \"/content/drive/MyDrive/soypod_crop_counting/soypod_crop_counting_{}.txt\".format(img_type)\n","    img_list = [name.split(',') for name in open(img_list_path).read().splitlines()]\n","    #\n","    # save the number of detected pods\n","    csv_name = (args.vis_dir).split(\"/\")[-1]\n","    loss_csv = open(os.path.join(args.vis_dir, '{}_seed_counting_{}.csv'.format(csv_name, img_type)), 'w+')\n","\n","    for img_path_ij, _ in img_list:\n","        print(\"image path = {}\".format(img_path_ij))\n","        #\n","        img_name = img_path_ij.split(\"/\")[-1]\n","        #\n","        img_name_pred_before = img_name.replace(\".png\", '_pred_bf.png')\n","        img_name_pred_after = img_name.replace(\".png\", '_pred_af.png')\n","        # load the images\n","        img_0 = cv2.imread(img_path_ij)\n","        # this is only for drawing points\n","        img_raw = cv2.cvtColor(img_0, cv2.COLOR_BGR2RGB)\n","        # pre-proccessing\n","        img = transform(img_raw)\n","        ##\n","        # round the size\n","        height, width = img_0.shape[:2]\n","        # get the new input image size suitable for VGG16 net\n","        new_width = (width // 128 +1) * 128\n","        new_height = (height // 128 + 1)* 128\n","        ##\n","        img_in = torch.zeros((3, new_height, new_width))\n","        img_in[:,:height,:width] = img\n","        ##\n","        ##\n","        if new_height > 1700:\n","            print(\"Large file\")\n","            # for out out image\n","            img_draw = (np.ones((new_height, new_width, 3))*255).astype(np.uint8)\n","            img_draw[:height, :width,:] = img_raw\n","            #\n","            img_to_draw_before = img_draw.copy() \n","            img_to_draw_after = img_draw.copy() \n","            #\n","            new_width_hf = int(new_width/2)\n","            new_height_hf = int(new_height/2)\n","            for shi in range(1,3):\n","                print(\"the {} half\".format(shi))\n","                # prepare the output image\n","                img_raw_shi = img_draw[new_height_hf*(shi-1):new_height_hf*shi,:,:]\n","                #\n","                samples = img_in[:,new_height_hf*(shi-1):new_height_hf*shi,:].unsqueeze(0).to(device)\n","                # run inference\n","                outputs = model(samples)\n","\n","                outputs_scores = torch.nn.functional.softmax(outputs['pred_logits'], -1)[:, :, 1][0]\n","\n","                outputs_points = outputs['pred_points'][0]\n","\n","                # filter the predictions\n","                # 0.5 is used by default\n","                points = outputs_points[outputs_scores > threshold].detach().cpu().numpy()#.tolist()\n","                if points.shape[0]< 10000 and points.shape[0] != 0:\n","                    #print(\"doing clustering\")\n","                    cutoff = 500/points.shape[0]\n","                    if cutoff<20:\n","                        cutoff = 20\n","                    components = nx.connected_components(\n","                        nx.from_edgelist(\n","                            (i, j) for i, js in enumerate(\n","                                spatial.KDTree(points).query_ball_point(points, cutoff)\n","                            )\n","                            for j in js\n","                        )\n","                    )\n","\n","                    clusters = {j: i for i, js in enumerate(components) for j in js}\n","\n","                    # reorganize the points to the order of clusters \n","                    points_reo = np.zeros(points.shape)\n","                    i = 0\n","                    for key in clusters.keys():\n","                        #print(key)\n","                        points_reo[i,:] = points[key,:]\n","                        i+=1\n","                    # points_n has the same order as clusters\n","                    res = [clusters[key] for key in clusters.keys()]\n","                    res_n = np.array(res).reshape(-1,1)\n","\n","                    points_n = []\n","                    for i in np.unique(res_n):\n","                        tmp = points_reo[np.where(res_n[:,0] == i)]\n","                        points_n.append( [np.mean(tmp[:,0]), np.mean(tmp[:,1])])\n","                else:\n","                    points_n = points.tolist()\n","                #\n","                if shi ==1:\n","                    points_bf_sum = np.array(points)\n","                    points_af_sum = np.array(points_n)\n","                    print(\"points_af_sum {}\".format(points_af_sum.shape))\n","                else:\n","                    points_bf_sum = np.concatenate((points_bf_sum, np.array(points)), 0)\n","                    points_af_sum = np.concatenate((points_af_sum, np.array(points_n)), 0)\n","                \n","                # draw the predictions\n","                alpha = 0.5\n","                #. before \n","                size = 6\n","                img_to_draw_before_in = cv2.cvtColor(np.array(img_raw_shi), cv2.COLOR_RGB2BGR)\n","                img_to_draw_before_in_n = img_to_draw_before_in.copy()\n","                for p in points:\n","                    img_to_draw_before_in_n = cv2.circle(img_to_draw_before_in_n, (int(p[0]), int(p[1])), size, (0, 0, 255), -1)\n","                img_to_draw_before_in_nn = cv2.addWeighted(img_to_draw_before_in_n, alpha, img_to_draw_before_in, 1 - alpha, 0)\n","                # save the visualized image\n","                img_to_draw_before[new_height_hf*(shi-1):new_height_hf*shi,:,:] = img_to_draw_before_in_nn\n","                #. after  \n","                #size = 6\n","                img_to_draw_after_in = cv2.cvtColor(np.array(img_raw_shi), cv2.COLOR_RGB2BGR)\n","                img_to_draw_after_in_n = img_to_draw_after_in.copy()\n","                for p in points_n:\n","                    img_to_draw_after_in_n = cv2.circle(img_to_draw_after_in_n, (int(p[0]), int(p[1])), size, (0, 0, 255), -1)\n","                #\n","                img_to_draw_after_in_nn = cv2.addWeighted(img_to_draw_after_in_n, alpha, img_to_draw_after_in, 1 - alpha, 0)\n","                img_to_draw_after[new_height_hf*(shi-1):new_height_hf*shi,:,:] = img_to_draw_after_in_nn\n","            #\n","            cv2.imwrite(os.path.join(args.vis_dir, img_name_pred_before), img_to_draw_before[:height,:width,:])\n","            # save the visualized image\n","            cv2.imwrite(os.path.join(args.vis_dir, img_name_pred_after), img_to_draw_after[:height,:width,:])\n","            #\n","            #predict_cnt = int((outputs_scores > threshold).sum())\n","            points_bf_sum = points_bf_sum.tolist()\n","            points_af_sum = points_af_sum.tolist()\n","            predict_cnt_before = len(points_bf_sum)\n","            predict_cnt_after = len(points_af_sum)\n","\n","        else:\n","            samples = img_in.unsqueeze(0).to(device)\n","            # run inference\n","            outputs = model(samples)\n","\n","            outputs_scores = torch.nn.functional.softmax(outputs['pred_logits'], -1)[:, :, 1][0]\n","\n","            outputs_points = outputs['pred_points'][0]\n","\n","            # filter the predictions\n","            # 0.5 is used by default\n","            points = outputs_points[outputs_scores > threshold].detach().cpu().numpy()#.tolist()\n","            if points.shape[0]< 10000 and points.shape[0] != 0:\n","                #print(\"doing clustering\")\n","                cutoff = 500/points.shape[0]\n","                if cutoff<20:\n","                    cutoff = 20\n","                components = nx.connected_components(\n","                    nx.from_edgelist(\n","                        (i, j) for i, js in enumerate(\n","                            spatial.KDTree(points).query_ball_point(points, cutoff)\n","                        )\n","                        for j in js\n","                    )\n","                )\n","\n","                clusters = {j: i for i, js in enumerate(components) for j in js}\n","\n","                # reorganize the points to the order of clusters \n","                points_reo = np.zeros(points.shape)\n","                i = 0\n","                for key in clusters.keys():\n","                    #print(key)\n","                    points_reo[i,:] = points[key,:]\n","                    i+=1\n","                # points_n has the same order as clusters\n","                res = [clusters[key] for key in clusters.keys()]\n","                res_n = np.array(res).reshape(-1,1)\n","\n","                points_n = []\n","                for i in np.unique(res_n):\n","                    tmp = points_reo[np.where(res_n[:,0] == i)]\n","                    points_n.append( [np.mean(tmp[:,0]), np.mean(tmp[:,1])])\n","            else:\n","                points_n = points.tolist()\n","            # calculate the distance and find the center of the too close points\n","\n","            #predict_cnt = int((outputs_scores > threshold).sum())\n","            predict_cnt_before = len(points)\n","            predict_cnt_after = len(points_n)\n","            #\n","            print(\"Number of seeds before = {}\".format(predict_cnt_before))\n","            print(\"Number of seeds after = {}\".format(predict_cnt_after))\n","            # draw the predictions\n","            alpha = 0.5\n","            #. before \n","            size = 6\n","            #\n","            img_to_draw_before = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR)\n","            img_to_draw_before_in_x = img_to_draw_before.copy()\n","            for p in points:\n","                img_to_draw_before_in_x = cv2.circle(img_to_draw_before_in_x, (int(p[0]), int(p[1])), size, (0, 0, 255), -1)\n","            img_to_draw_before_in_xx = cv2.addWeighted(img_to_draw_before_in_x, alpha, img_to_draw_before, 1 - alpha, 0)\n","            # save the visualized image\n","            cv2.imwrite(os.path.join(args.vis_dir, img_name_pred_before), img_to_draw_before_in_xx[:height,:width,:])\n","            #. after  \n","            #size = 6\n","            img_to_draw_after = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR)\n","            img_to_draw_after_in_x = img_to_draw_after.copy()\n","            for p in points_n:\n","                img_to_draw_after_in_x = cv2.circle(img_to_draw_after_in_x, (int(p[0]), int(p[1])), size, (0, 0, 255), -1)\n","            img_to_draw_after_in_xx = cv2.addWeighted(img_to_draw_after_in_x, alpha, img_to_draw_after, 1 - alpha, 0)\n","            # save the visualized image\n","            cv2.imwrite(os.path.join(args.vis_dir, img_name_pred_after), img_to_draw_after_in_xx[:height,:width,:])\n","            #\n","        print(\"Number of seeds before = {}\".format(predict_cnt_before))\n","        print(\"Number of seeds after = {}\".format(predict_cnt_after))\n","        # save the detected pod number\n","        loss_csv.write('{},{},{},{}\\n'.format(img_name_pred_before.split(\".\")[0], predict_cnt_before, img_name_pred_after.split(\".\")[0], predict_cnt_after))\n","        loss_csv.flush()  \n","    loss_csv.close\n"],"metadata":{"id":"SDlL7Xrn8I0O"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["#####On big infield images"],"metadata":{"id":"-e9cOLjZ8qZV"}},{"cell_type":"code","source":["# threshold for evaluation\n","threshold = 0.5 #args.threshold\n","# to apply post processing or not: if Filter = 10000, the filter is applied\n","\n","#\n","# load trained model\n","if args.resume is not None:\n","    checkpoint = torch.load(args.resume, map_location='cpu')\n","    model.load_state_dict(checkpoint['model'])\n","#\n","model.eval()\n","#\n","step = 0\n","epoch = 0\n","###\n","# create the pre-processing transform\n","transform = standard_transforms.Compose([\n","    standard_transforms.ToTensor(), \n","    standard_transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n","])\n","\n","# set your image path here\n","folder_path = \"/content/drive/MyDrive/soybean_1118selection_annotated/\"\n","folder_sub = os.listdir(folder_path)\n","#\n","# save the number of detected pods\n","csv_name = (args.vis_dir).split(\"/\")[-1]\n","loss_csv = open(os.path.join(args.vis_dir, '{}_seed_counting_{}.csv'.format(csv_name, 'final')), 'w+')\n","#\n","img_ct = 0\n","for sub in range(len(folder_sub)):\n","    #\n","    folder_sub_s = folder_sub[sub]\n","    #\n","    folder_path_s = os.path.join(folder_path, folder_sub_s)\n","    #\n","    img_list = glob.glob(os.path.join(folder_path_s,'*.JPG'))#[:3]\n","    #\n","    for img_path_ij in img_list:\n","        print(\"image path = {}\".format(img_path_ij))\n","        #\n","        img_name = img_path_ij.split(\"/\")[-1]\n","        #\n","        img_name_pred_before = img_name.replace(\".png\", '_pred_bf.png')\n","        img_name_pred_after = img_name.replace(\".png\", '_pred_af.png')\n","        # load the images\n","        img_00 = cv2.imread(img_path_ij)\n","        img_0 = cv2.resize(img_00, (int(img_00.shape[1]/2), int(img_00.shape[0]/2)), interpolation = cv2.INTER_AREA)\n","        # this is only for drawing points\n","        #img_raw = Image.fromarray(cv2.cvtColor(img_0, cv2.COLOR_BGR2RGB))\n","        img_raw = cv2.cvtColor(img_0, cv2.COLOR_BGR2RGB)\n","        # pre-proccessing\n","        img = transform(img_raw)\n","        ##\n","        # round the size\n","        height, width = img_0.shape[:2]\n","        # get the new input image size suitable for VGG16 net\n","        new_width = (width // 128 +1) * 128\n","        new_height = (height // 128 + 1)* 128\n","        ##\n","        img_in = torch.zeros((3, new_height, new_width))\n","        img_in[:,:height,:width] = img\n","\n","        # divide the image into two to save memory\n","        print(\"new_height = {}\".format(new_height))\n","        # for out out image\n","        img_draw = (np.ones((new_height, new_width, 3))*255).astype(np.uint8)\n","        img_draw[:height, :width,:] = img_raw\n","        #\n","        img_to_draw_before = img_draw.copy() \n","        img_to_draw_after = img_draw.copy() \n","        #\n","        new_width_hf = int(new_width/4)\n","        new_height_hf = int(new_height/4)\n","        for shi in range(1,5):\n","            print(\"the {} half\".format(shi))\n","            # prepare the output image\n","            img_raw_shi = img_draw[:,new_width_hf*(shi-1):new_width_hf*shi,:]\n","            #\n","            samples = img_in[:,:,new_width_hf*(shi-1):new_width_hf*shi].unsqueeze(0).to(device)\n","            # run inference\n","            outputs = model(samples)\n","\n","            outputs_scores = torch.nn.functional.softmax(outputs['pred_logits'], -1)[:, :, 1][0]\n","\n","            outputs_points = outputs['pred_points'][0]\n","\n","            # filter the predictions\n","            # 0.5 is used by default\n","            points = outputs_points[outputs_scores > threshold].detach().cpu().numpy()#.tolist()\n","            if points.shape[0]< 100000 and points.shape[0] != 0:\n","                #print(\"doing clustering\")\n","                cutoff = 500/points.shape[0]\n","                if cutoff<10:\n","                    cutoff = 10\n","                components = nx.connected_components(\n","                    nx.from_edgelist(\n","                        (i, j) for i, js in enumerate(\n","                            spatial.KDTree(points).query_ball_point(points, cutoff)\n","                        )\n","                        for j in js\n","                    )\n","                )\n","\n","                clusters = {j: i for i, js in enumerate(components) for j in js}\n","\n","                # reorganize the points to the order of clusters \n","                points_reo = np.zeros(points.shape)\n","                i = 0\n","                for key in clusters.keys():\n","                    #print(key)\n","                    points_reo[i,:] = points[key,:]\n","                    i+=1\n","                # points_n has the same order as clusters\n","                res = [clusters[key] for key in clusters.keys()]\n","                res_n = np.array(res).reshape(-1,1)\n","\n","                points_n = []\n","                for i in np.unique(res_n):\n","                    tmp = points_reo[np.where(res_n[:,0] == i)]\n","                    points_n.append( [np.mean(tmp[:,0]), np.mean(tmp[:,1])])\n","            else:\n","                points_n = points.tolist()\n","            #\n","            if points_n:\n","                    if shi ==1:\n","                        points_bf_sum = np.array(points)\n","                        points_af_sum = np.array(points_n)\n","                        print(\"points_af_sum {}\".format(points_af_sum.shape))\n","                    else:\n","                        points_bf_sum = np.concatenate((points_bf_sum, np.array(points)), 0)\n","                        points_af_sum = np.concatenate((points_af_sum, np.array(points_n)), 0)\n","            # draw the predictions\n","            alpha = 0.6\n","            #. before \n","            size = 5\n","            img_to_draw_before_in = cv2.cvtColor(np.array(img_raw_shi), cv2.COLOR_RGB2BGR)\n","            img_to_draw_before_in_n = img_to_draw_before_in.copy()\n","            for p in points:\n","                img_to_draw_before_in_n = cv2.circle(img_to_draw_before_in_n, (int(p[0]), int(p[1])), size, (0, 0, 255), -1)\n","            img_to_draw_before_in_nn = cv2.addWeighted(img_to_draw_before_in_n, alpha, img_to_draw_before_in, 1 - alpha, 0)\n","            # save the visualized image\n","            img_to_draw_before[:,new_width_hf*(shi-1):new_width_hf*shi,:] = img_to_draw_before_in_nn\n","            #. after  \n","            #size = 4\n","            img_to_draw_after_in = cv2.cvtColor(np.array(img_raw_shi), cv2.COLOR_RGB2BGR)\n","            img_to_draw_after_in_n = img_to_draw_after_in.copy()\n","            for p in points_n:\n","                img_to_draw_after_in_n = cv2.circle(img_to_draw_after_in_n, (int(p[0]), int(p[1])), size, (0, 0, 255), -1)\n","            #\n","            img_to_draw_after_in_nn = cv2.addWeighted(img_to_draw_after_in_n, alpha, img_to_draw_after_in, 1 - alpha, 0)\n","            img_to_draw_after[:,new_width_hf*(shi-1):new_width_hf*shi,:] = img_to_draw_after_in_nn\n","        #\n","        cv2.imwrite(os.path.join(args.vis_dir, img_name_pred_before), img_to_draw_before[:height,:width,:])\n","        # save the visualized image\n","        cv2.imwrite(os.path.join(args.vis_dir, img_name_pred_after), img_to_draw_after[:height,:width,:])\n","        #\n","        #predict_cnt = int((outputs_scores > threshold).sum())\n","        points_bf_sum = points_bf_sum.tolist()\n","        points_af_sum = points_af_sum.tolist()\n","        predict_cnt_before = len(points_bf_sum)\n","        predict_cnt_after = len(points_af_sum)\n","        #\n","        print(\"Number of seeds before = {}\".format(predict_cnt_before))\n","        print(\"Number of seeds after = {}\".format(predict_cnt_after))\n","        # save the detected pod number\n","        loss_csv.write('{},{},{},{}\\n'.format(img_name_pred_before.split(\".\")[0], predict_cnt_before, img_name_pred_after.split(\".\")[0], predict_cnt_after))\n","        loss_csv.flush()  \n","        #\n","        print(\"saving {}th image {}\".format(img_ct, img_path_ij))\n","        img_ct +=1\n","loss_csv.close\n"],"metadata":{"id":"e0OPDItL8I27"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":[],"metadata":{"id":"xGm84VVk8I5n"},"execution_count":null,"outputs":[]}],"metadata":{"accelerator":"GPU","colab":{"collapsed_sections":[],"machine_shape":"hm","provenance":[{"file_id":"1qkGFGfTqkhKWtzLwPJrmTPQ3V-C6Kr2z","timestamp":1667883619593}]},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"},"widgets":{"application/vnd.jupyter.widget-state+json":{"80ec627ed5f3417ebecbbcfe767338aa":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_ec2555b45230458896dbeb40f581d433","IPY_MODEL_db3be6617bf34876be88384c1a6e9597","IPY_MODEL_fa532929d8774047abdc95b5af3296a9"],"layout":"IPY_MODEL_4be575664ce449ddb5cbef4377147e38"}},"ec2555b45230458896dbeb40f581d433":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_ceed4556f6474d2bbaeba192fa0cb1e5","placeholder":"​","style":"IPY_MODEL_690f8ed4dcaa4f95b415366319ab5944","value":"100%"}},"db3be6617bf34876be88384c1a6e9597":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_44a8014a98b3499692332bf9838f151a","max":553433881,"min":0,"orientation":"horizontal","style":"IPY_MODEL_ce00405a13e945aa89616c7947053ae3","value":553433881}},"fa532929d8774047abdc95b5af3296a9":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_675bf8a98fbd4208b1f5537c8fe7ecd4","placeholder":"​","style":"IPY_MODEL_4c31f05ca00541a2b018c6a8bfc393e6","value":" 528M/528M [00:08&lt;00:00, 24.6MB/s]"}},"4be575664ce449ddb5cbef4377147e38":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"ceed4556f6474d2bbaeba192fa0cb1e5":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"690f8ed4dcaa4f95b415366319ab5944":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"44a8014a98b3499692332bf9838f151a":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"ce00405a13e945aa89616c7947053ae3":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"675bf8a98fbd4208b1f5537c8fe7ecd4":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"4c31f05ca00541a2b018c6a8bfc393e6":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}}}}},"nbformat":4,"nbformat_minor":0}