<gh_stars>0
from __future__ import division, print_function, absolute_import
import re
import glob
import os.path as osp

from torch.utils.data import Dataset
from senseTk.common import *
from senseTk.functions import LAP_Matching
import numpy as np
import torch

from torchreid.utils import read_image
from torchreid.data.transforms import build_transforms


class MOT16(Dataset):
    """MOT16 format tracking dataset.

    Reference:
        - Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.
        - Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.

    URL: `<https://github.com/layumi/DukeMTMC-reID_evaluation>`_
    
    Dataset statistics:
        - identities: 1404 (train + query).
        - images:16522 (train) + 2228 (query) + 17661 (gallery).
        - cameras: 8.
    """
    dataset_dir = 'MOT16'

    # dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip'

    def __init__(
        self,
        root='',
        mode='train',
        test_in_train_format=False,
        meta_file='',
        step=30,
        h_len=3,
        height=256,
        width=128,
        loc_std=1,
        **kwargs
    ):
        self.root = osp.abspath(osp.expanduser(root))
        self.dataset_dir = osp.join(self.root, self.dataset_dir)
        self.train_im_dir = osp.join(self.dataset_dir, 'imgs/train')
        self.val_im_dir = osp.join(self.dataset_dir, 'imgs/train')
        self.test_im_dir = osp.join(self.dataset_dir, 'imgs/test')
        self.train_dt_dir = osp.join(self.dataset_dir, 'dts/train')
        self.val_dt_dir = osp.join(self.dataset_dir, 'dts/train')
        self.test_dt_dir = osp.join(self.dataset_dir, 'dts/test')
        self.train_gt_dir = osp.join(self.dataset_dir, 'gts/train')
        self.val_gt_dir = osp.join(self.dataset_dir, 'gts/train')

        # required_files = [
        #     self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir
        # ]
        # self.check_before_run(required_files)

        self.test_mode = mode
        self.meta_file = meta_file
        self.test_in_train_format = test_in_train_format

        self.step = step
        self.h_len = h_len
        self.height = height
        self.width = width
        self.loc_std = loc_std
        self.transform_tr, self.transform_te = build_transforms(
            self.height,
            self.width,
        )

        assert self.test_mode in [
            'train', 'val', 'test'
        ], 'mode must be one of [train, val or test]'
        self.sequences = []
        self.data = self._process()

        self._cache_flag = False


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

    def _match_gt(self, dt, gt):
        for f in gt.frameRange():
            d = dt[f]
            g = gt[f]
            m, l, r = LAP_Matching(d, g, Det.iou)
            for i, j in m:
                if d[i].iou(g[j]) < 0.55:
                    continue
                tmp = d[i]
                dt.delete(tmp)
                tmp.uid = g[j].uid
                dt.append_data(tmp)

    def _process(self):
        if self.meta_file == '':
            self.meta_file = osp.join(
                self.dataset_dir, self.test_mode + '.list.txt'
            )
        data = []
        with open(self.meta_file) as fd:
            lines = fd.readlines()
        im_mappings = {
            'test': self.test_im_dir,
            'val': self.val_im_dir,
            'train': self.train_im_dir,
        }
        dt_mappings = {
            'test': self.test_dt_dir,
            'val': self.val_dt_dir,
            'train': self.train_dt_dir,
        }
        gt_mappings = {
            'val': self.val_gt_dir,
            'train': self.train_gt_dir,
        }
        for one in lines:
            one = one.strip()
            seq_imdir = osp.join(im_mappings[self.test_mode], one)
            vid = VideoClipReader(seq_imdir)
            seq_dt = osp.join(dt_mappings[self.test_mode], one + '.txt')
            dt = TrackSet(seq_dt, formatter='fr.i id.i x1 y1 w h cf -1 -1 -1')
            if self.test_mode == 'train' or self.test_in_train_format:
                seq_gt = osp.join(gt_mappings[self.test_mode], one + '.txt')
                if not osp.exists(seq_gt):
                    seq_gt = seq_gt[:-4]
                gt = TrackSet(
                    seq_gt,
                    formatter='fr.i id.i x1 y1 w h st.i -1 -1 -1',
                    # filter=lambda d: d.st == 1,
                )
                self._match_gt(dt, gt)
                #for gid in gt
                n = len(vid) - 1
                self.sequences.append([])
                for i in range(1, n, self.step):
                    be = i
                    en = min(i + self.step, n)
                    cur = en
                    im_path = osp.join(
                        vid.backend.i_root,
                        vid.backend.fmt % (vid.backend.start + en - 1)
                    )
                    for d in dt[cur]:
                        if d.uid < 0:
                            continue
                        ref_ims = []
                        ref_dts = []
                        ref_las = []
                        data_ = {
                            'cur_im': im_path,
                            'cur_dt': [d.x1, d.y1, d.x2 + 1, d.y2 + 1],
                            'ref_im': ref_ims,
                            'ref_dt': ref_dts,
                            'ref_la': ref_las
                        }
                        cnt = 0
                        for j in range(self.h_len):
                            tar = en - 1 - j * self.step // self.h_len
                            tar = max(tar, 1)
                            tmp = dt(d.uid)[tar]
                            pre_im_path = osp.join(
                                vid.backend.i_root, vid.backend.fmt %
                                (vid.backend.start + tar - 1)
                            )
                            if len(tmp):
                                pre_d = tmp[0]
                            else:
                                pre_d = None
                            use_neg = np.random.choice(2)
                            if pre_d is None or use_neg:
                                tmp_n = len(dt[tar])
                                if tmp_n <= 0:
                                    pre_d = d.copy()
                                    pre_d.uid = -1
                                else:
                                    ind = np.random.choice(tmp_n)
                                    pre_d = dt[tar][ind]
                            if pre_d.uid == d.uid:
                                cnt += 1
                            ref_ims.append(pre_im_path)
                            ref_dts.append(
                                [
                                    pre_d.x1, pre_d.y1, pre_d.x2 + 1,
                                    pre_d.y2 + 1
                                ]
                            )
                            ref_las.append(pre_d.uid == d.uid)
                        if cnt >= int(self.h_len / 3. * 2):
                            label = 1
                        else:
                            label = 0
                        data_['label'] = label
                        self.sequences[-1].append(len(data))
                        data.append(data_)
            else:
                n = len(vid)
                for i in range(n):
                    im_path = osp.join(
                        vid.backend.i_root,
                        vid.backend.fmt % (vid.backend.start + i)
                    )
                    bboxes = []
                    for d in dt[vid.backend.start + i]:
                        bboxes.append([d.x1, d.y1, d.x2 + 1, d.y2 + 1, d.conf])
                    data_ = {
                        'img': im_path,
                        'dets': bboxes,
                        'fr': i,
                        'total': n,
                        'seq': one
                    }
                    data.append(data_)
        return data

    def _cache_init(self):
        if not self._cache_flag:
            # self._cache_ims = dequeue(self.h_len * 5)
            self._cache_flag = True

    def _read_im(self, p):
        # for i in self._cache_ims:
        #     if i[0] == p:
        #         return i[1]
        im = read_image(p)
        # self._cache_ims.append((p, im))
        return im

    def _crop(self, im, d):
        x1, y1, x2, y2 = map(int, d[:4])
        crop = im.crop((x1, y1, x2, y2))
        return crop

    def _process_im(self, path, dets):
        self._cache_init()
        single = False
        if isinstance(path, str):
            path = [path]
            dets = dets.reshape(1, -1)
            single = True
        ret = []
        for p, d in zip(path, dets):
            im = self._read_im(p)
            ret.append(self._crop(im, d))
        if single:
            return ret[0]
        else:
            return ret

    def _pipeline_im(self, d):
        if self.test_mode == 'train' or self.test_in_train_format:
            return self.transform_tr(d)
        else:
            return self.transform_te(d)

    def pipeline(self, data):
        for k in data:
            if 'im' in k:
                if isinstance(data[k], list):
                    for i in range(len(data[k])):
                        data[k][i] = self._pipeline_im(data[k][i])
                    data[k] = torch.stack(data[k])
                else:
                    data[k] = self._pipeline_im(data[k])
                    data[k] = data[k].unsqueeze(0)
            if 'dt' in k:
                data[k] = torch.from_numpy(data[k]).float()
                data[k] /= self.loc_std
            if 'dets' in k:
                data['dt'] = torch.from_numpy(data[k]).float()
                data['dt'] /= self.loc_std
        return data

    def __getitem__(self, ind):
        data = self.data[ind]
        ret = {}
        if self.test_mode != 'train' and not self.test_in_train_format:
            ret['dets'] = np.array(data['dets'])
            ret['im'] = self._process_im(
                [data['img']] * len(data['dets']), data['dets']
            )
            for k in ['fr', 'total', 'seq']:
                ret[k] = data[k]
        else:
            ret['cur_dt'] = np.array(data['cur_dt'])
            ret['ref_dt'] = np.array(data['ref_dt'])
            ret['cur_im'] = self._process_im(data['cur_im'], ret['cur_dt'])
            ret['ref_im'] = self._process_im(data['ref_im'], ret['ref_dt'])
            ret['label'] = data['label']
            ret['ref_la'] = np.array(data['ref_la'])
        ret = self.pipeline(ret)
        return ret


class HIE20(MOT16):

    dataset_dir = 'HIE20'

    pass
from __future__ import division, print_function, absolute_import
import re
import glob
import os.path as osp

from torch.utils.data import Dataset
from senseTk.common import *
from senseTk.functions import LAP_Matching
import numpy as np
import torch

from torchreid.utils import read_image
from torchreid.data.transforms import build_transforms


class MOT16(Dataset):
    """MOT16 format tracking dataset.

    Reference:
        - Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.
        - Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.

    URL: `<https://github.com/layumi/DukeMTMC-reID_evaluation>`_
    
    Dataset statistics:
        - identities: 1404 (train + query).
        - images:16522 (train) + 2228 (query) + 17661 (gallery).
        - cameras: 8.
    """
    dataset_dir = 'MOT16'

    # dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip'

    def __init__(
        self,
        root='',
        mode='train',
        test_in_train_format=False,
        meta_file='',
        step=30,
        h_len=3,
        height=256,
        width=128,
        loc_std=1,
        **kwargs
    ):
        self.root = osp.abspath(osp.expanduser(root))
        self.dataset_dir = osp.join(self.root, self.dataset_dir)
        self.train_im_dir = osp.join(self.dataset_dir, 'imgs/train')
        self.val_im_dir = osp.join(self.dataset_dir, 'imgs/train')
        self.test_im_dir = osp.join(self.dataset_dir, 'imgs/test')
        self.train_dt_dir = osp.join(self.dataset_dir, 'dts/train')
        self.val_dt_dir = osp.join(self.dataset_dir, 'dts/train')
        self.test_dt_dir = osp.join(self.dataset_dir, 'dts/test')
        self.train_gt_dir = osp.join(self.dataset_dir, 'gts/train')
        self.val_gt_dir = osp.join(self.dataset_dir, 'gts/train')

        # required_files = [
        #     self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir
        # ]
        # self.check_before_run(required_files)

        self.test_mode = mode
        self.meta_file = meta_file
        self.test_in_train_format = test_in_train_format

        self.step = step
        self.h_len = h_len
        self.height = height
        self.width = width
        self.loc_std = loc_std
        self.transform_tr, self.transform_te = build_transforms(
            self.height,
            self.width,
        )

        assert self.test_mode in [
            'train', 'val', 'test'
        ], 'mode must be one of [train, val or test]'
        self.sequences = []
        self.data = self._process()

        self._cache_flag = False


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

    def _match_gt(self, dt, gt):
        for f in gt.frameRange():
            d = dt[f]
            g = gt[f]
            m, l, r = LAP_Matching(d, g, Det.iou)
            for i, j in m:
                if d[i].iou(g[j]) < 0.55:
                    continue
                tmp = d[i]
                dt.delete(tmp)
                tmp.uid = g[j].uid
                dt.append_data(tmp)

    def _process(self):
        if self.meta_file == '':
            self.meta_file = osp.join(
                self.dataset_dir, self.test_mode + '.list.txt'
            )
        data = []
        with open(self.meta_file) as fd:
            lines = fd.readlines()
        im_mappings = {
            'test': self.test_im_dir,
            'val': self.val_im_dir,
            'train': self.train_im_dir,
        }
        dt_mappings = {
            'test': self.test_dt_dir,
            'val': self.val_dt_dir,
            'train': self.train_dt_dir,
        }
        gt_mappings = {
            'val': self.val_gt_dir,
            'train': self.train_gt_dir,
        }
        for one in lines:
            one = one.strip()
            seq_imdir = osp.join(im_mappings[self.test_mode], one)
            vid = VideoClipReader(seq_imdir)
            seq_dt = osp.join(dt_mappings[self.test_mode], one + '.txt')
            dt = TrackSet(seq_dt, formatter='fr.i id.i x1 y1 w h cf -1 -1 -1')
            if self.test_mode == 'train' or self.test_in_train_format:
                seq_gt = osp.join(gt_mappings[self.test_mode], one + '.txt')
                if not osp.exists(seq_gt):
                    seq_gt = seq_gt[:-4]
                gt = TrackSet(
                    seq_gt,
                    formatter='fr.i id.i x1 y1 w h st.i -1 -1 -1',
                    # filter=lambda d: d.st == 1,
                )
                self._match_gt(dt, gt)
                #for gid in gt
                n = len(vid) - 1
                self.sequences.append([])
                for i in range(1, n, self.step):
                    be = i
                    en = min(i + self.step, n)
                    cur = en
                    im_path = osp.join(
                        vid.backend.i_root,
                        vid.backend.fmt % (vid.backend.start + en - 1)
                    )
                    for d in dt[cur]:
                        if d.uid < 0:
                            continue
                        ref_ims = []
                        ref_dts = []
                        ref_las = []
                        data_ = {
                            'cur_im': im_path,
                            'cur_dt': [d.x1, d.y1, d.x2 + 1, d.y2 + 1],
                            'ref_im': ref_ims,
                            'ref_dt': ref_dts,
                            'ref_la': ref_las
                        }
                        cnt = 0
                        for j in range(self.h_len):
                            tar = en - 1 - j * self.step // self.h_len
                            tar = max(tar, 1)
                            tmp = dt(d.uid)[tar]
                            pre_im_path = osp.join(
                                vid.backend.i_root, vid.backend.fmt %
                                (vid.backend.start + tar - 1)
                            )
                            if len(tmp):
                                pre_d = tmp[0]
                            else:
                                pre_d = None
                            use_neg = np.random.choice(2)
                            if pre_d is None or use_neg:
                                tmp_n = len(dt[tar])
                                if tmp_n <= 0:
                                    pre_d = d.copy()
                                    pre_d.uid = -1
                                else:
                                    ind = np.random.choice(tmp_n)
                                    pre_d = dt[tar][ind]
                            if pre_d.uid == d.uid:
                                cnt += 1
                            ref_ims.append(pre_im_path)
                            ref_dts.append(
                                [
                                    pre_d.x1, pre_d.y1, pre_d.x2 + 1,
                                    pre_d.y2 + 1
                                ]
                            )
                            ref_las.append(pre_d.uid == d.uid)
                        if cnt >= int(self.h_len / 3. * 2):
                            label = 1
                        else:
                            label = 0
                        data_['label'] = label
                        self.sequences[-1].append(len(data))
                        data.append(data_)
            else:
                n = len(vid)
                for i in range(n):
                    im_path = osp.join(
                        vid.backend.i_root,
                        vid.backend.fmt % (vid.backend.start + i)
                    )
                    bboxes = []
                    for d in dt[vid.backend.start + i]:
                        bboxes.append([d.x1, d.y1, d.x2 + 1, d.y2 + 1, d.conf])
                    data_ = {
                        'img': im_path,
                        'dets': bboxes,
                        'fr': i,
                        'total': n,
                        'seq': one
                    }
                    data.append(data_)
        return data

    def _cache_init(self):
        if not self._cache_flag:
            # self._cache_ims = dequeue(self.h_len * 5)
            self._cache_flag = True

    def _read_im(self, p):
        # for i in self._cache_ims:
        #     if i[0] == p:
        #         return i[1]
        im = read_image(p)
        # self._cache_ims.append((p, im))
        return im

    def _crop(self, im, d):
        x1, y1, x2, y2 = map(int, d[:4])
        crop = im.crop((x1, y1, x2, y2))
        return crop

    def _process_im(self, path, dets):
        self._cache_init()
        single = False
        if isinstance(path, str):
            path = [path]
            dets = dets.reshape(1, -1)
            single = True
        ret = []
        for p, d in zip(path, dets):
            im = self._read_im(p)
            ret.append(self._crop(im, d))
        if single:
            return ret[0]
        else:
            return ret

    def _pipeline_im(self, d):
        if self.test_mode == 'train' or self.test_in_train_format:
            return self.transform_tr(d)
        else:
            return self.transform_te(d)

    def pipeline(self, data):
        for k in data:
            if 'im' in k:
                if isinstance(data[k], list):
                    for i in range(len(data[k])):
                        data[k][i] = self._pipeline_im(data[k][i])
                    data[k] = torch.stack(data[k])
                else:
                    data[k] = self._pipeline_im(data[k])
                    data[k] = data[k].unsqueeze(0)
            if 'dt' in k:
                data[k] = torch.from_numpy(data[k]).float()
                data[k] /= self.loc_std
            if 'dets' in k:
                data['dt'] = torch.from_numpy(data[k]).float()
                data['dt'] /= self.loc_std
        return data

    def __getitem__(self, ind):
        data = self.data[ind]
        ret = {}
        if self.test_mode != 'train' and not self.test_in_train_format:
            ret['dets'] = np.array(data['dets'])
            ret['im'] = self._process_im(
                [data['img']] * len(data['dets']), data['dets']
            )
            for k in ['fr', 'total', 'seq']:
                ret[k] = data[k]
        else:
            ret['cur_dt'] = np.array(data['cur_dt'])
            ret['ref_dt'] = np.array(data['ref_dt'])
            ret['cur_im'] = self._process_im(data['cur_im'], ret['cur_dt'])
            ret['ref_im'] = self._process_im(data['ref_im'], ret['ref_dt'])
            ret['label'] = data['label']
            ret['ref_la'] = np.array(data['ref_la'])
        ret = self.pipeline(ret)
        return ret


class HIE20(MOT16):

    dataset_dir = 'HIE20'

    pass
MOT16 format tracking dataset.      Reference:         - Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.         - Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.      URL: `<https://github.com/layumi/DukeMTMC-reID_evaluation>`_          Dataset statistics:         - identities: 1404 (train + query).         - images:16522 (train) + 2228 (query) + 17661 (gallery).         - cameras: 8. # dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip' # required_files = [ #     self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir # ] # self.check_before_run(required_files) # filter=lambda d: d.st == 1, #for gid in gt # self._cache_ims = dequeue(self.h_len * 5) # for i in self._cache_ims: #     if i[0] == p: #         return i[1] # self._cache_ims.append((p, im))