# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

"""
This file contains components with some default boilerplate logic user may need
in training / testing. They will not work for everyone, but many users may find them useful.

The behavior of functions/classes in this file is subject to change,
since they are meant to represent the "common default behavior" people need in their projects.
"""

import argparse
import logging
import os
import sys
from collections import OrderedDict
import torch
from fvcore.common.file_io import PathManager
from fvcore.nn.precise_bn import get_bn_modules
from torch.nn.parallel import DistributedDataParallel

import detectron2.data.transforms as T
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import (
    MetadataCatalog,
    build_detection_test_loader,
    build_detection_train_loader,
)
from detectron2.evaluation import (
    DatasetEvaluator,
    inference_on_dataset,
    print_csv_format,
    verify_results,
)
from detectron2.modeling import build_model
from detectron2.solver import build_lr_scheduler, build_optimizer
from detectron2.utils import comm
from detectron2.utils.collect_env import collect_env_info
from detectron2.utils.env import seed_all_rng
from detectron2.utils.events import CommonMetricPrinter, JSONWriter, TensorboardXWriter
from detectron2.utils.logger import setup_logger
from detectron2.structures import BoxMode

from torch.utils.data import Dataset, DataLoader
from torch.autograd import Variable

import cv2
from tqdm import tqdm
import json

def get_pedestrain_dict(image_list, json_list):
    dataset_dicts = []
    
    for i,path in tqdm(enumerate(image_list)):
        filename = path
        #img = cv2.imread(path)
        # height, width = cv2.imread(filename).shape[:2]
        record = {}
        record['file_name'] = filename
        #record['file_img'] = img
        record['image_id'] = i #path.split('/')[-1][:-5]
        #id is like 000000 or 000001
        record['height']= 480
        record['width']= 640
        
        #for i in data_list[1] to get bbox and category
        objs = []
        
        people = json_list[i]
        with open(people) as p:
            json_context = json.load(p)
            for person in json_context:
                boxes = list(map(float, person['pos']))
                obj = {
                    "bbox": boxes,
                    "bbox_mode": BoxMode.XYWH_ABS,
                    #"segmentation": [poly], To draw a line, along to ballon
                    "category_id": 0,
                    "iscrowd": 0
                }
                objs.append(obj)
            record["annotations"] = objs
        dataset_dicts.append(record)
    return dataset_dicts #list of dicts


class inference_Dataset(Dataset):
    def __init__(self, train_d, dic, intputFormat, aug):
        self.filenames = train_d
        
        self.dic = dic
        self.input_format = intputFormat
        
        self.aug = aug

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, idx):
        original_image = cv2.imread(self.filenames[idx])
        
        if self.input_format == "RGB":
            # whether the model expects BGR inputs or RGB
            original_image = original_image[:, :, ::-1]
        height, width = original_image.shape[:2]
        image = self.aug.get_transform(original_image).apply_image(original_image)
        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
        
        #img = Image.open(self.filenames[idx])  # PIL image
        #image = self.transform(img)
        return image, height, width


class DefaultBatchPredictor:
    """
    Create a simple end-to-end predictor with the given config that runs on
    single device for a single input image.

    Compared to using the model directly, this class does the following additions:

    1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
    2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
    3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
    4. Take one input image and produce a single output, instead of a batch.

    If you'd like to do anything more fancy, please refer to its source code
    as examples to build and use the model manually.

    Attributes:
        metadata (Metadata): the metadata of the underlying dataset, obtained from
            cfg.DATASETS.TEST.

    Examples:
    ::
        pred = DefaultPredictor(cfg)
        inputs = cv2.imread("input.jpg")
        outputs = pred(inputs)
    """

    def __init__(self, cfg, togpu=False, modelPath=None):
        self.cfg = cfg.clone()  # cfg can be modified by model
#         self.cfg.MODEL.WEIGHTS = modelPath     
        self.model = build_model(self.cfg)
        self.model.eval()

        checkpointer = DetectionCheckpointer(self.model)
        
        if modelPath==None:
            checkpointer.load(self.cfg.MODEL.WEIGHTS)
            print("Load weight from : ", self.cfg.MODEL.WEIGHTS)
        else:
            checkpointer.load(modelPath)
            print("Load weight from : ", modelPath)

        self.aug = T.ResizeShortestEdge(
            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
        )

        self.input_format = cfg.INPUT.FORMAT
        assert self.input_format in ["RGB", "BGR"], self.input_format
        
    def create_batch_loader(self, dataPath, datajson, batch=24, workers=4):
        
        self.dic  = get_pedestrain_dict(dataPath, datajson)
        self.ba_dataloader = DataLoader(
                                        inference_Dataset(dataPath, self.dic , intputFormat = self.input_format, aug = self.aug),
                                        batch_size=batch,
                                        num_workers=workers
                                        )
        return self.ba_dataloader
        #print(self.ba_dataloader.__len__())
    
    def batchPredictor(self):
        result=[]
        result_d=[]
        with torch.no_grad():
            for x,h,w in tqdm(self.ba_dataloader):
                inputs = []  
                for i in range(len(x)):
                    inputs.append({"image":x[i], "height": h[i], "width": w[i]})
                outputs = self.model(inputs)
                
                #print(dic)
                #print(">>>>>>>" ,len(dic))
                for i in outputs:
                    result.append(i)
                    #print(i)
                #break
        return result, self.dic

    def __call__(self, original_image):
        """
        Args:
            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).

        Returns:
            predictions (dict):
                the output of the model for one image only.
                See :doc:`/tutorials/models` for details about the format.
        """
        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
            # Apply pre-processing to image.
            if self.input_format == "RGB":
                # whether the model expects BGR inputs or RGB
                original_image = original_image[:, :, ::-1]
            height, width = original_image.shape[:2]
            image = self.aug.get_transform(original_image).apply_image(original_image)
            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))

            inputs = {"image": image, "height": height, "width": width}
            predictions = self.model([inputs])[0]
            return predictions
    

# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

"""
This file contains components with some default boilerplate logic user may need
in training / testing. They will not work for everyone, but many users may find them useful.

The behavior of functions/classes in this file is subject to change,
since they are meant to represent the "common default behavior" people need in their projects.
"""

import argparse
import logging
import os
import sys
from collections import OrderedDict
import torch
from fvcore.common.file_io import PathManager
from fvcore.nn.precise_bn import get_bn_modules
from torch.nn.parallel import DistributedDataParallel

import detectron2.data.transforms as T
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import (
    MetadataCatalog,
    build_detection_test_loader,
    build_detection_train_loader,
)
from detectron2.evaluation import (
    DatasetEvaluator,
    inference_on_dataset,
    print_csv_format,
    verify_results,
)
from detectron2.modeling import build_model
from detectron2.solver import build_lr_scheduler, build_optimizer
from detectron2.utils import comm
from detectron2.utils.collect_env import collect_env_info
from detectron2.utils.env import seed_all_rng
from detectron2.utils.events import CommonMetricPrinter, JSONWriter, TensorboardXWriter
from detectron2.utils.logger import setup_logger
from detectron2.structures import BoxMode

from torch.utils.data import Dataset, DataLoader
from torch.autograd import Variable

import cv2
from tqdm import tqdm
import json

def get_pedestrain_dict(image_list, json_list):
    dataset_dicts = []
    
    for i,path in tqdm(enumerate(image_list)):
        filename = path
        #img = cv2.imread(path)
        # height, width = cv2.imread(filename).shape[:2]
        record = {}
        record['file_name'] = filename
        #record['file_img'] = img
        record['image_id'] = i #path.split('/')[-1][:-5]
        #id is like 000000 or 000001
        record['height']= 480
        record['width']= 640
        
        #for i in data_list[1] to get bbox and category
        objs = []
        
        people = json_list[i]
        with open(people) as p:
            json_context = json.load(p)
            for person in json_context:
                boxes = list(map(float, person['pos']))
                obj = {
                    "bbox": boxes,
                    "bbox_mode": BoxMode.XYWH_ABS,
                    #"segmentation": [poly], To draw a line, along to ballon
                    "category_id": 0,
                    "iscrowd": 0
                }
                objs.append(obj)
            record["annotations"] = objs
        dataset_dicts.append(record)
    return dataset_dicts #list of dicts


class inference_Dataset(Dataset):
    def __init__(self, train_d, dic, intputFormat, aug):
        self.filenames = train_d
        
        self.dic = dic
        self.input_format = intputFormat
        
        self.aug = aug

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, idx):
        original_image = cv2.imread(self.filenames[idx])
        
        if self.input_format == "RGB":
            # whether the model expects BGR inputs or RGB
            original_image = original_image[:, :, ::-1]
        height, width = original_image.shape[:2]
        image = self.aug.get_transform(original_image).apply_image(original_image)
        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
        
        #img = Image.open(self.filenames[idx])  # PIL image
        #image = self.transform(img)
        return image, height, width


class DefaultBatchPredictor:
    """
    Create a simple end-to-end predictor with the given config that runs on
    single device for a single input image.

    Compared to using the model directly, this class does the following additions:

    1. Load checkpoint from `cfg.MODEL.WEIGHTS`.
    2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.
    3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.
    4. Take one input image and produce a single output, instead of a batch.

    If you'd like to do anything more fancy, please refer to its source code
    as examples to build and use the model manually.

    Attributes:
        metadata (Metadata): the metadata of the underlying dataset, obtained from
            cfg.DATASETS.TEST.

    Examples:
    ::
        pred = DefaultPredictor(cfg)
        inputs = cv2.imread("input.jpg")
        outputs = pred(inputs)
    """

    def __init__(self, cfg, togpu=False, modelPath=None):
        self.cfg = cfg.clone()  # cfg can be modified by model
#         self.cfg.MODEL.WEIGHTS = modelPath     
        self.model = build_model(self.cfg)
        self.model.eval()

        checkpointer = DetectionCheckpointer(self.model)
        
        if modelPath==None:
            checkpointer.load(self.cfg.MODEL.WEIGHTS)
            print("Load weight from : ", self.cfg.MODEL.WEIGHTS)
        else:
            checkpointer.load(modelPath)
            print("Load weight from : ", modelPath)

        self.aug = T.ResizeShortestEdge(
            [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
        )

        self.input_format = cfg.INPUT.FORMAT
        assert self.input_format in ["RGB", "BGR"], self.input_format
        
    def create_batch_loader(self, dataPath, datajson, batch=24, workers=4):
        
        self.dic  = get_pedestrain_dict(dataPath, datajson)
        self.ba_dataloader = DataLoader(
                                        inference_Dataset(dataPath, self.dic , intputFormat = self.input_format, aug = self.aug),
                                        batch_size=batch,
                                        num_workers=workers
                                        )
        return self.ba_dataloader
        #print(self.ba_dataloader.__len__())
    
    def batchPredictor(self):
        result=[]
        result_d=[]
        with torch.no_grad():
            for x,h,w in tqdm(self.ba_dataloader):
                inputs = []  
                for i in range(len(x)):
                    inputs.append({"image":x[i], "height": h[i], "width": w[i]})
                outputs = self.model(inputs)
                
                #print(dic)
                #print(">>>>>>>" ,len(dic))
                for i in outputs:
                    result.append(i)
                    #print(i)
                #break
        return result, self.dic

    def __call__(self, original_image):
        """
        Args:
            original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).

        Returns:
            predictions (dict):
                the output of the model for one image only.
                See :doc:`/tutorials/models` for details about the format.
        """
        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
            # Apply pre-processing to image.
            if self.input_format == "RGB":
                # whether the model expects BGR inputs or RGB
                original_image = original_image[:, :, ::-1]
            height, width = original_image.shape[:2]
            image = self.aug.get_transform(original_image).apply_image(original_image)
            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))

            inputs = {"image": image, "height": height, "width": width}
            predictions = self.model([inputs])[0]
            return predictions
    

# -*- coding: utf-8 -*- # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved This file contains components with some default boilerplate logic user may need in training / testing. They will not work for everyone, but many users may find them useful.  The behavior of functions/classes in this file is subject to change, since they are meant to represent the "common default behavior" people need in their projects. #img = cv2.imread(path) # height, width = cv2.imread(filename).shape[:2] #record['file_img'] = img #path.split('/')[-1][:-5] #id is like 000000 or 000001 #for i in data_list[1] to get bbox and category #"segmentation": [poly], To draw a line, along to ballon #list of dicts # whether the model expects BGR inputs or RGB #img = Image.open(self.filenames[idx])  # PIL image #image = self.transform(img) Create a simple end-to-end predictor with the given config that runs on     single device for a single input image.      Compared to using the model directly, this class does the following additions:      1. Load checkpoint from `cfg.MODEL.WEIGHTS`.     2. Always take BGR image as the input and apply conversion defined by `cfg.INPUT.FORMAT`.     3. Apply resizing defined by `cfg.INPUT.{MIN,MAX}_SIZE_TEST`.     4. Take one input image and produce a single output, instead of a batch.      If you'd like to do anything more fancy, please refer to its source code     as examples to build and use the model manually.      Attributes:         metadata (Metadata): the metadata of the underlying dataset, obtained from             cfg.DATASETS.TEST.      Examples:     ::         pred = DefaultPredictor(cfg)         inputs = cv2.imread("input.jpg")         outputs = pred(inputs) # cfg can be modified by model #         self.cfg.MODEL.WEIGHTS = modelPath #print(self.ba_dataloader.__len__()) #print(dic) #print(">>>>>>>" ,len(dic)) #print(i) #break Args:             original_image (np.ndarray): an image of shape (H, W, C) (in BGR order).          Returns:             predictions (dict):                 the output of the model for one image only.                 See :doc:`/tutorials/models` for details about the format. # https://github.com/sphinx-doc/sphinx/issues/4258 # Apply pre-processing to image. # whether the model expects BGR inputs or RGB