import time
import json
import base64
from io import BytesIO

import torch.nn.functional as F
from torch.autograd import Variable
from flask import Flask, request, jsonify

from model.locator import Crowd_locator
import misc.transforms as own_transforms
from misc.utils import *
from collections import OrderedDict
import server_config

cfg = server_config.get_cfg_defaults()
server_config.show_cfg()

GPU_ID = cfg.GPU_ID
os.environ["CUDA_VISIBLE_DEVICES"] = GPU_ID
torch.backends.cudnn.benchmark = True

netName = cfg.NET
model_path = cfg.WEIGHT

mean_std = cfg.MEAN_STD

img_transform = standard_transforms.Compose([
    standard_transforms.ToTensor(),
    standard_transforms.Normalize(*mean_std)
])

restore = standard_transforms.Compose([
    own_transforms.DeNormalize(*mean_std),
    standard_transforms.ToPILImage()
])


def write_json(path: str, data):
    with open(path, 'w') as f:
        json.dump(data, f)


def get_boxInfo_from_Binar_map(Binar_numpy, min_area=3):
    Binar_numpy = Binar_numpy.squeeze().astype(np.uint8)
    assert Binar_numpy.ndim == 2
    cnt, labels, stats, centroids = cv2.connectedComponentsWithStats(Binar_numpy, connectivity=4)  # centriod (w,h)
    # cnt: 返回值是连通区域的数量。
    # labels: labels是一个与image一样大小的矩形（labels.shape = image.shape），其中每一个连通区域会有一个唯一标识，标识从0开始。
    # stats会包含5个参数分别为x,y,h,w,s(s=连通区域的像素个数)
    # centroids: 返回的是连通区域的质心。
    boxes = stats[1:, :]
    points = centroids[1:, :]
    index = (boxes[:, 4] >= min_area)
    boxes = boxes[index]
    points = points[index]
    pre_data = {'num': len(points), 'points': points}
    return pre_data, boxes


def load_weights(net, model_path: str):
    pretrain_state_dict = torch.load(model_path)

    if len(GPU_ID.split(',')) > 1:
        net.load_state_dict(pretrain_state_dict)

    else:
        new_state_dict = OrderedDict()
        for k, v in pretrain_state_dict.items():
            name = k.replace('module.', '')
            new_state_dict[name] = v
        net.load_state_dict(new_state_dict)


# init model
net = Crowd_locator(netName, GPU_ID, pretrained=True)
net.cuda()
print('load net model')
load_weights(net, model_path)
net.eval()

app = Flask(__name__)


def _base64_to_pil(base64_data):
    img = None
    for i, data in enumerate(base64_data):
        decode_data = base64.b64decode(data)
        img_data = BytesIO(decode_data)
        img = Image.open(img_data)
    return img


@app.route('/crowd', methods=['POST'])
def index():
    params = request.json if request.method == "POST" else request.args
    try:
        img = _base64_to_pil(params['img'])
        img = img.convert('RGB') if img.mode == 'L' else img
        pred_data, boxes = prediction(img)
        data = {
            "code": 200,
            "timestamp": int(time.time()),
            "num": int(pred_data['num']),
            "points": pred_data['points'].tolist(),
            "boxes": boxes.tolist()
        }
    except Exception as e:
        print(e)
        data = {
            "code": 500,
            "message": str(e)
        }
    return jsonify(data)


def run():
    app.config['JSON_AS_ASCII'] = False
    app.run(host=cfg.IP, port=int(cfg.PORT), use_reloader=False)


def prediction(img):
    img = img_transform(img)[None, :, :, :]
    slice_h, slice_w = 512, 1024

    with torch.no_grad():
        img = Variable(img).cuda()
        b, c, h, w = img.shape

        crop_imgs, crop_dots, crop_masks = [], [], []

        if h * w < slice_h * 2 * slice_w * 2 and h % 16 == 0 and w % 16 == 0:
            [pred_threshold, pred_map, __] = [i.cpu() for i in net(img, mask_gt=None, mode='val')]

        else:
            if h % 16 != 0:
                pad_dims = (0, 0, 0, 16 - h % 16)
                h = (h // 16 + 1) * 16
                img = F.pad(img, pad_dims, "constant")

            if w % 16 != 0:
                pad_dims = (0, 16 - w % 16, 0, 0)
                w = (w // 16 + 1) * 16
                img = F.pad(img, pad_dims, "constant")

            for i in range(0, h, slice_h):
                h_start, h_end = max(min(h - slice_h, i), 0), min(h, i + slice_h)
                for j in range(0, w, slice_w):
                    w_start, w_end = max(min(w - slice_w, j), 0), min(w, j + slice_w)
                    crop_imgs.append(img[:, :, h_start:h_end, w_start:w_end])
                    mask = torch.zeros(1, 1, img.size(2), img.size(3)).cpu()
                    mask[:, :, h_start:h_end, w_start:w_end].fill_(1.0)
                    crop_masks.append(mask)

            crop_imgs, crop_masks = torch.cat(crop_imgs, dim=0), torch.cat(crop_masks, dim=0)

            # forward may need repeatng
            crop_preds, crop_thresholds = [], []
            nz, period = crop_imgs.size(0), 4
            for i in range(0, nz, period):
                [crop_threshold, crop_pred, __] = [i.cpu() for i in
                                                   net(crop_imgs[i:min(nz, i + period)], mask_gt=None, mode='val')]
                crop_preds.append(crop_pred)
                crop_thresholds.append(crop_threshold)

            crop_preds = torch.cat(crop_preds, dim=0)
            crop_thresholds = torch.cat(crop_thresholds, dim=0)

            # splice them to the original size
            idx = 0
            pred_map = torch.zeros(b, 1, h, w).cpu()
            pred_threshold = torch.zeros(b, 1, h, w).cpu().float()
            for i in range(0, h, slice_h):
                h_start, h_end = max(min(h - slice_h, i), 0), min(h, i + slice_h)
                for j in range(0, w, slice_w):
                    w_start, w_end = max(min(w - slice_w, j), 0), min(w, j + slice_w)
                    pred_map[:, :, h_start:h_end, w_start:w_end] += crop_preds[idx]
                    pred_threshold[:, :, h_start:h_end, w_start:w_end] += crop_thresholds[idx]
                    idx += 1
            mask = crop_masks.sum(dim=0)
            pred_map = (pred_map / mask)
            pred_threshold = (pred_threshold / mask)

        a = torch.ones_like(pred_map)
        b = torch.zeros_like(pred_map)
        binar_map = torch.where(pred_map >= pred_threshold, a, b)

        pred_data, boxes = get_boxInfo_from_Binar_map(binar_map.cpu().numpy())

    return pred_data, boxes


if __name__ == '__main__':
    run()

import time
import json
import base64
from io import BytesIO

import torch.nn.functional as F
from torch.autograd import Variable
from flask import Flask, request, jsonify

from model.locator import Crowd_locator
import misc.transforms as own_transforms
from misc.utils import *
from collections import OrderedDict
import server_config

cfg = server_config.get_cfg_defaults()
server_config.show_cfg()

GPU_ID = cfg.GPU_ID
os.environ["CUDA_VISIBLE_DEVICES"] = GPU_ID
torch.backends.cudnn.benchmark = True

netName = cfg.NET
model_path = cfg.WEIGHT

mean_std = cfg.MEAN_STD

img_transform = standard_transforms.Compose([
    standard_transforms.ToTensor(),
    standard_transforms.Normalize(*mean_std)
])

restore = standard_transforms.Compose([
    own_transforms.DeNormalize(*mean_std),
    standard_transforms.ToPILImage()
])


def write_json(path: str, data):
    with open(path, 'w') as f:
        json.dump(data, f)


def get_boxInfo_from_Binar_map(Binar_numpy, min_area=3):
    Binar_numpy = Binar_numpy.squeeze().astype(np.uint8)
    assert Binar_numpy.ndim == 2
    cnt, labels, stats, centroids = cv2.connectedComponentsWithStats(Binar_numpy, connectivity=4)  # centriod (w,h)
    # cnt: 返回值是连通区域的数量。
    # labels: labels是一个与image一样大小的矩形（labels.shape = image.shape），其中每一个连通区域会有一个唯一标识，标识从0开始。
    # stats会包含5个参数分别为x,y,h,w,s(s=连通区域的像素个数)
    # centroids: 返回的是连通区域的质心。
    boxes = stats[1:, :]
    points = centroids[1:, :]
    index = (boxes[:, 4] >= min_area)
    boxes = boxes[index]
    points = points[index]
    pre_data = {'num': len(points), 'points': points}
    return pre_data, boxes


def load_weights(net, model_path: str):
    pretrain_state_dict = torch.load(model_path)

    if len(GPU_ID.split(',')) > 1:
        net.load_state_dict(pretrain_state_dict)

    else:
        new_state_dict = OrderedDict()
        for k, v in pretrain_state_dict.items():
            name = k.replace('module.', '')
            new_state_dict[name] = v
        net.load_state_dict(new_state_dict)


# init model
net = Crowd_locator(netName, GPU_ID, pretrained=True)
net.cuda()
print('load net model')
load_weights(net, model_path)
net.eval()

app = Flask(__name__)


def _base64_to_pil(base64_data):
    img = None
    for i, data in enumerate(base64_data):
        decode_data = base64.b64decode(data)
        img_data = BytesIO(decode_data)
        img = Image.open(img_data)
    return img


@app.route('/crowd', methods=['POST'])
def index():
    params = request.json if request.method == "POST" else request.args
    try:
        img = _base64_to_pil(params['img'])
        img = img.convert('RGB') if img.mode == 'L' else img
        pred_data, boxes = prediction(img)
        data = {
            "code": 200,
            "timestamp": int(time.time()),
            "num": int(pred_data['num']),
            "points": pred_data['points'].tolist(),
            "boxes": boxes.tolist()
        }
    except Exception as e:
        print(e)
        data = {
            "code": 500,
            "message": str(e)
        }
    return jsonify(data)


def run():
    app.config['JSON_AS_ASCII'] = False
    app.run(host=cfg.IP, port=int(cfg.PORT), use_reloader=False)


def prediction(img):
    img = img_transform(img)[None, :, :, :]
    slice_h, slice_w = 512, 1024

    with torch.no_grad():
        img = Variable(img).cuda()
        b, c, h, w = img.shape

        crop_imgs, crop_dots, crop_masks = [], [], []

        if h * w < slice_h * 2 * slice_w * 2 and h % 16 == 0 and w % 16 == 0:
            [pred_threshold, pred_map, __] = [i.cpu() for i in net(img, mask_gt=None, mode='val')]

        else:
            if h % 16 != 0:
                pad_dims = (0, 0, 0, 16 - h % 16)
                h = (h // 16 + 1) * 16
                img = F.pad(img, pad_dims, "constant")

            if w % 16 != 0:
                pad_dims = (0, 16 - w % 16, 0, 0)
                w = (w // 16 + 1) * 16
                img = F.pad(img, pad_dims, "constant")

            for i in range(0, h, slice_h):
                h_start, h_end = max(min(h - slice_h, i), 0), min(h, i + slice_h)
                for j in range(0, w, slice_w):
                    w_start, w_end = max(min(w - slice_w, j), 0), min(w, j + slice_w)
                    crop_imgs.append(img[:, :, h_start:h_end, w_start:w_end])
                    mask = torch.zeros(1, 1, img.size(2), img.size(3)).cpu()
                    mask[:, :, h_start:h_end, w_start:w_end].fill_(1.0)
                    crop_masks.append(mask)

            crop_imgs, crop_masks = torch.cat(crop_imgs, dim=0), torch.cat(crop_masks, dim=0)

            # forward may need repeatng
            crop_preds, crop_thresholds = [], []
            nz, period = crop_imgs.size(0), 4
            for i in range(0, nz, period):
                [crop_threshold, crop_pred, __] = [i.cpu() for i in
                                                   net(crop_imgs[i:min(nz, i + period)], mask_gt=None, mode='val')]
                crop_preds.append(crop_pred)
                crop_thresholds.append(crop_threshold)

            crop_preds = torch.cat(crop_preds, dim=0)
            crop_thresholds = torch.cat(crop_thresholds, dim=0)

            # splice them to the original size
            idx = 0
            pred_map = torch.zeros(b, 1, h, w).cpu()
            pred_threshold = torch.zeros(b, 1, h, w).cpu().float()
            for i in range(0, h, slice_h):
                h_start, h_end = max(min(h - slice_h, i), 0), min(h, i + slice_h)
                for j in range(0, w, slice_w):
                    w_start, w_end = max(min(w - slice_w, j), 0), min(w, j + slice_w)
                    pred_map[:, :, h_start:h_end, w_start:w_end] += crop_preds[idx]
                    pred_threshold[:, :, h_start:h_end, w_start:w_end] += crop_thresholds[idx]
                    idx += 1
            mask = crop_masks.sum(dim=0)
            pred_map = (pred_map / mask)
            pred_threshold = (pred_threshold / mask)

        a = torch.ones_like(pred_map)
        b = torch.zeros_like(pred_map)
        binar_map = torch.where(pred_map >= pred_threshold, a, b)

        pred_data, boxes = get_boxInfo_from_Binar_map(binar_map.cpu().numpy())

    return pred_data, boxes


if __name__ == '__main__':
    run()

# centriod (w,h) # cnt: 返回值是连通区域的数量。 # labels: labels是一个与image一样大小的矩形（labels.shape = image.shape），其中每一个连通区域会有一个唯一标识，标识从0开始。 # stats会包含5个参数分别为x,y,h,w,s(s=连通区域的像素个数) # centroids: 返回的是连通区域的质心。 # init model # forward may need repeatng # splice them to the original size