# Copyright 2021 Tencent

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
import os
import numpy as np
import torch
import warnings
import random
import matplotlib.pyplot as plt
import gradio as gr
import torchvision.transforms as standard_transforms
from torch.utils.data import DataLoader
from torch.utils.data import Dataset

warnings.filterwarnings('ignore')

# define the GPU id to be used
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

class data(Dataset):
    def __init__(self, img, transform=None):
        self.image = img
        self.transform = transform

    def __len__(self):
        return 1000

    def __getitem__(self, x):
        # open image here as PIL / numpy
        image = self.image
        image = image.convert('RGB')
        if self.transform is not None:
            image = self.transform(image)
      
        image = torch.Tensor(image)
        return image

def loading_data(img):
    # the augumentations
    transform = standard_transforms.Compose([
        standard_transforms.ToTensor(), standard_transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                                      std=[0.229, 0.224, 0.225]),
    ])
    # dcreate  the dataset
    test_set = data(img=img, transform=transform)
    test_loader = DataLoader(test_set, batch_size=1, num_workers=4, shuffle=False, drop_last=False)

    return test_loader


def predict(img):
    """the main process of inference"""
    test_loader = loading_data(img)
    model = SASNet().cuda()
    model_path = "./SHHA.pth"
    # load the trained model
    model.load_state_dict(torch.load(model_path))
    print('successfully load model from', model_path)

    with torch.no_grad():
        model.eval()

        for vi, data in enumerate(test_loader, 0):
          img = data
          img = img.cuda()
          pred_map = model(img)
          pred_map = pred_map.data.cpu().numpy()
          for i_img in range(pred_map.shape[0]):
            pred_cnt = np.sum(pred_map[i_img]) / 1000

            den_map = np.squeeze(pred_map[i_img])
            fig = plt.figure(frameon=False)
            ax = plt.Axes(fig, [0., 0., 1., 1.])
            ax.set_axis_off()
            fig.add_axes(ax)
            ax.imshow(den_map, aspect='auto')
            
            return pred_cnt, fig


with gr.Blocks() as demo:
    gr.Markdown("""
    # Crowd Counting based on SASNet

    We implemented a image crowd counting model with VGG16 following the paper of Song et. al (2021).
    
    ## Abstract
In this paper, we address the large scale variation problem in crowd counting by taking full advantage of the multi-scale feature representations in a multi-level network. We implement such an idea by keeping the counting error of a patch as small as possible with a proper feature level selection strategy, since a specific feature level tends to perform better for a certain range of scales. However, without scale annotations, it is sub-optimal and error-prone to manually assign the predictions for heads of different scales to specific feature levels. Therefore, we propose a Scale-Adaptive Selection Network (SASNet), which automatically learns the internal correspondence between the scales and the feature levels. Instead of directly using the predictions from the most appropriate feature level as the final estimation, our SASNet also considers the predictions from other feature levels via weighted average, which helps to mitigate the gap between discrete feature levels and continuous scale variation. Since the heads in a local patch share roughly a same scale, we conduct the adaptive selection strategy in a patch-wise style. However, pixels within a patch contribute different counting errors due to the various difficulty degrees of learning. Thus, we further propose a Pyramid Region Awareness Loss (PRA Loss) to recursively select the most hard sub-regions within a patch until reaching the pixel level. With awareness of whether the parent patch is over-estimated or under-estimated, the fine-grained optimization with the PRA Loss for these region-aware hard pixels helps to alleviate the inconsistency problem between training target and evaluation metric. The state-of-the-art results on four datasets demonstrate the superiority of our approach. 

The code will be available at: https://github.com/TencentYoutuResearch/CrowdCounting-SASNet. 

    ## References
    Song, Q., Wang, C., Wang, Y., Tai, Y., Wang, C., Li, J., … Ma, J. (2021). To Choose or to Fuse? Scale Selection for Crowd Counting. The Thirty-Fifth AAAI Conference on Artificial Intelligence (AAAI-21).
    """)
    image_button = gr.Button("Count the Crowd!")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil")
            gr.Examples(["IMG_1.jpg", "IMG_2.jpg", "IMG_3.jpg"], image_input)
        with gr.Column():
            image_output = gr.Plot()
        with gr.Column():
            text_output = gr.Label()
            

    image_button.click(predict, inputs=image_input, outputs=[text_output, image_output])

demo.launch(debug = True)