#!/usr/bin/env python
import time
import os
import numpy as np

from chroma import generator
from chroma import gpu
from chroma import event
from chroma import itertoolset
from chroma.tools import profile_if_possible

import pycuda.driver as cuda

from timeit import default_timer as timer

def pick_seed():
    """Returns a seed for a random number generator selected using
    a mixture of the current time and the current process ID."""
    return int(time.time()) ^ (os.getpid() << 16)

class Simulation(object):
    def __init__(self, detector, seed=None, cuda_device=None, particle_tracking=False, photon_tracking=False,
                 geant4_processes=4, nthreads_per_block=64, max_blocks=1024):
        self.detector = detector

        self.nthreads_per_block = nthreads_per_block
        self.max_blocks = max_blocks
        self.photon_tracking = photon_tracking

        if seed is None:
            self.seed = pick_seed()
        else:
            self.seed = seed

        # We have three generators to seed: numpy.random, GEANT4, and CURAND.
        # The latter two are done below.
        np.random.seed(self.seed)

        if geant4_processes > 0:
            self.photon_generator = generator.photon.G4ParallelGenerator(geant4_processes, detector.detector_material, base_seed=self.seed, tracking=particle_tracking)
        else:
            self.photon_generator = None

        self.context = gpu.create_cuda_context(cuda_device)

        if hasattr(detector, 'num_channels'):
            self.gpu_geometry = gpu.GPUDetector(detector)
            self.gpu_daq = gpu.GPUDaq(self.gpu_geometry)
            self.gpu_pdf = gpu.GPUPDF()
            self.gpu_pdf_kernel = gpu.GPUKernelPDF()
        else:
            self.gpu_geometry = gpu.GPUGeometry(detector)

        self.rng_states = gpu.get_rng_states(self.nthreads_per_block*self.max_blocks, seed=self.seed)

        self.pdf_config = None

    def _simulate_batch(self,batch_events,keep_photons_beg=False,keep_photons_end=False,keep_hits=True,keep_flat_hits=True,run_daq=False, max_steps=100, verbose=False):
        '''Assumes batch_events is a list of Event objects with photons_beg having evidx set to the index in the array.

           Yields the fully formed events. Do not call directly.'''

        t_start = timer()

        #Idea: allocate memory on gpu and copy photons into it, instead of concatenating on CPU?
        batch_photons = event.Photons.join([ev.photons_beg for ev in batch_events])
        batch_bounds = np.cumsum(np.concatenate([[0],[len(ev.photons_beg) for ev in batch_events]]))

        #This copy to gpu has a _lot_ of overhead, want 100k photons at least, hence batches
        #Assume triangles, and weights are unimportant to copy to GPU
        t_copy_start = timer()
        gpu_photons = gpu.GPUPhotons(batch_photons,copy_triangles=False,copy_weights=False)
        t_copy_end = timer()
        if verbose:
            print('GPU copy took %0.2f s' % (t_copy_end-t_copy_start))

        t_prop_start = timer()
        tracking = gpu_photons.propagate(self.gpu_geometry, self.rng_states,
                              nthreads_per_block=self.nthreads_per_block,
                              max_blocks=self.max_blocks,
                              max_steps=max_steps,track=self.photon_tracking)

        t_prop_end = timer()
        if verbose:
            print('GPU propagate took %0.2f s' % (t_prop_end-t_prop_start))

        t_end = timer()
        if verbose:
            print('Batch took %0.2f s' % (t_end-t_start))

        if keep_photons_end:
            batch_photons_end = gpu_photons.get()

        if hasattr(self.detector, 'num_channels') and (keep_hits or keep_flat_hits):
            batch_hits = gpu_photons.get_flat_hits(self.gpu_geometry)

        for i,(batch_ev,(start_photon,end_photon)) in enumerate(zip(batch_events,zip(batch_bounds[:-1],batch_bounds[1:]))):

            if not keep_photons_beg:
                batch_ev.photons_beg = None

            if self.photon_tracking:
                step_photon_ids,step_photons = tracking
                nphotons = end_photon-start_photon
                photon_tracks = [[] for i in range(nphotons)]
                for step_ids,step_photons in zip(step_photon_ids,step_photons):
                    mask = np.logical_and(step_ids >= start_photon,step_ids<end_photon)
                    if np.count_nonzero(mask) == 0:
                        break
                    photon_ids = step_ids[mask]-start_photon
                    photons = step_photons[mask]
                    #Indexing Photons with a scalar changes the internal array shapes...
                    any(photon_tracks[id].append(photons[i]) for i,id in enumerate(photon_ids))
                batch_ev.photon_tracks = [event.Photons.join(photons,concatenate=False) if len(photons) > 0 else event.Photons() for photons in photon_tracks]

            if keep_photons_end:
                batch_ev.photons_end = batch_photons_end[start_photon:end_photon]

            if hasattr(self.detector, 'num_channels') and (keep_hits or keep_flat_hits):
                ev_hits = batch_hits[batch_hits.evidx == i]
                if keep_hits:
                    #This is kind of expensive computationally, but keep_hits is for diagnostics
                    batch_ev.hits = { int(chan):ev_hits[ev_hits.channel == chan] for chan in np.unique(ev_hits.channel) }
                if keep_flat_hits:
                    batch_ev.flat_hits = ev_hits

            if hasattr(self, 'gpu_daq') and run_daq:
                #Must run DAQ per event, or design a much more complicated daq algorithm
                self.gpu_daq.begin_acquire()
                self.gpu_daq.acquire(gpu_photons, self.rng_states,
                                     start_photon=start_photon,
                                     nphotons=(end_photon-start_photon),
                                     nthreads_per_block=self.nthreads_per_block,
                                     max_blocks=self.max_blocks)
                gpu_channels = self.gpu_daq.end_acquire()
                batch_ev.channels = gpu_channels.get()

            yield batch_ev

    def simulate(self, iterable, keep_photons_beg=False, keep_photons_end=False,
                 keep_hits=True, keep_flat_hits=True, run_daq=False, max_steps=1000,
                 photons_per_batch=1000000):
        if isinstance(iterable, event.Photons):
            first_element, iterable = iterable, [iterable]
        else:
            first_element, iterable = itertoolset.peek(iterable)

        if isinstance(first_element, event.Event):
            iterable = self.photon_generator.generate_events(iterable)
        elif isinstance(first_element, event.Photons):
            iterable = (event.Event(photons_beg=x) for x in iterable)
        elif isinstance(first_element, event.Vertex):
            iterable = (event.Event(vertices=[vertex]) for vertex in iterable)
            iterable = self.photon_generator.generate_events(iterable)

        nphotons = 0
        batch_events = []

        for ev in iterable:

            ev.nphotons = len(ev.photons_beg)
            ev.photons_beg.evidx[:] = len(batch_events)

            nphotons += ev.nphotons
            batch_events.append(ev)

            #FIXME need an alternate implementation to split an event that is too large
            if nphotons >= photons_per_batch:
                yield from self._simulate_batch(batch_events,
                                                keep_photons_beg=keep_photons_beg,
                                                keep_photons_end=keep_photons_end,
                                                keep_hits=keep_hits,
                                                keep_flat_hits=keep_flat_hits,
                                                run_daq=run_daq, max_steps=max_steps)
                nphotons = 0
                batch_events = []

        if len(batch_events) != 0:
            yield from self._simulate_batch(batch_events,
                                            keep_photons_beg=keep_photons_beg,
                                            keep_photons_end=keep_photons_end,
                                            keep_hits=keep_hits,
                                            keep_flat_hits=keep_flat_hits,
                                            run_daq=run_daq, max_steps=max_steps)

    def create_pdf(self, iterable, tbins, trange, qbins, qrange, nreps=1):
        """Returns tuple: 1D array of channel hit counts, 3D array of
        (channel, time, charge) pdfs."""
        first_element, iterable = itertoolset.peek(iterable)

        if isinstance(first_element, event.Event):
            iterable = self.photon_generator.generate_events(iterable)

        pdf_config = (tbins, trange, qbins, qrange)
        if pdf_config != self.pdf_config:
            self.pdf_config = pdf_config
            self.gpu_pdf.setup_pdf(self.detector.num_channels(), tbins, trange,
                                   qbins, qrange)
        else:
            self.gpu_pdf.clear_pdf()

        if nreps > 1:
            iterable = itertoolset.repeating_iterator(iterable, nreps)

        for ev in iterable:
            gpu_photons = gpu.GPUPhotons(ev.photons_beg)
            gpu_photons.propagate(self.gpu_geometry, self.rng_states,
                                  nthreads_per_block=self.nthreads_per_block,
                                  max_blocks=self.max_blocks)
            self.gpu_daq.begin_acquire()
            self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks)
            gpu_channels = self.gpu_daq.end_acquire()
            self.gpu_pdf.add_hits_to_pdf(gpu_channels)

        return self.gpu_pdf.get_pdfs()

    @profile_if_possible
    def eval_pdf(self, event_channels, iterable, min_twidth, trange, min_qwidth, qrange, min_bin_content=100, nreps=1, ndaq=1, nscatter=1, time_only=True):
        """Returns tuple: 1D array of channel hit counts, 1D array of PDF
        probability densities."""
        ndaq_per_rep = 64
        ndaq_reps = ndaq // ndaq_per_rep
        gpu_daq = gpu.GPUDaq(self.gpu_geometry, ndaq=ndaq_per_rep)

        self.gpu_pdf.setup_pdf_eval(event_channels.hit,
                                    event_channels.t,
                                    event_channels.q,
                                    min_twidth,
                                    trange,
                                    min_qwidth,
                                    qrange,
                                    min_bin_content=min_bin_content,
                                    time_only=True)

        first_element, iterable = itertoolset.peek(iterable)

        if isinstance(first_element, event.Event):
            iterable = self.photon_generator.generate_events(iterable)
        elif isinstance(first_element, event.Photons):
            iterable = (event.Event(photons_beg=x) for x in iterable)

        for ev in iterable:
            gpu_photons_no_scatter = gpu.GPUPhotons(ev.photons_beg, ncopies=nreps)
            gpu_photons_scatter = gpu.GPUPhotons(ev.photons_beg, ncopies=nreps*nscatter)
            gpu_photons_no_scatter.propagate(self.gpu_geometry, self.rng_states,
                                             nthreads_per_block=self.nthreads_per_block,
                                             max_blocks=self.max_blocks,
                                             use_weights=True,
                                             scatter_first=-1,
                                             max_steps=10)
            gpu_photons_scatter.propagate(self.gpu_geometry, self.rng_states,
                                          nthreads_per_block=self.nthreads_per_block,
                                          max_blocks=self.max_blocks,
                                          use_weights=True,
                                          scatter_first=1,
                                          max_steps=5)
            nphotons = gpu_photons_no_scatter.true_nphotons # same for scatter
            for i in range(gpu_photons_no_scatter.ncopies):
                start_photon = i * nphotons
                gpu_photon_no_scatter_slice = gpu_photons_no_scatter.select(event.SURFACE_DETECT,
                                                                            start_photon=start_photon,
                                                                            nphotons=nphotons)
                gpu_photon_scatter_slices = [gpu_photons_scatter.select(event.SURFACE_DETECT,
                                                                        start_photon=(nscatter*i+j)*nphotons,
                                                                        nphotons=nphotons)
                                             for j in range(nscatter)]

                if len(gpu_photon_no_scatter_slice) == 0:
                    continue

                #weights = gpu_photon_slice.weights.get()
                #print 'weights', weights.min(), weights.max()
                for j in range(ndaq_reps):
                    gpu_daq.begin_acquire()
                    gpu_daq.acquire(gpu_photon_no_scatter_slice, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks)
                    for scatter_slice in gpu_photon_scatter_slices:
                        gpu_daq.acquire(scatter_slice, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks, weight=1.0/nscatter)
                    gpu_channels = gpu_daq.end_acquire()
                    self.gpu_pdf.accumulate_pdf_eval(gpu_channels, nthreads_per_block=ndaq_per_rep)

        return self.gpu_pdf.get_pdf_eval()

    def setup_kernel(self, event_channels, bandwidth_iterable,
                         trange, qrange,
                         nreps=1, ndaq=1, time_only=True, scale_factor=1.0):
        '''Call this before calling eval_pdf_kernel().  Sets up the
        event information and computes an appropriate kernel bandwidth'''
        nchannels = len(event_channels.hit)
        self.gpu_pdf_kernel.setup_moments(nchannels, trange, qrange,
                                          time_only=time_only)
        # Compute bandwidth
        first_element, bandwidth_iterable = itertoolset.peek(bandwidth_iterable)
        if isinstance(first_element, event.Event):
            bandwidth_iterable = \
                self.photon_generator.generate_events(bandwidth_iterable)
        for ev in bandwidth_iterable:
            gpu_photons = gpu.GPUPhotons(ev.photons_beg, ncopies=nreps)
            gpu_photons.propagate(self.gpu_geometry, self.rng_states,
                                  nthreads_per_block=self.nthreads_per_block,
                                  max_blocks=self.max_blocks)
            for gpu_photon_slice in gpu_photons.iterate_copies():
                for idaq in range(ndaq):
                    self.gpu_daq.begin_acquire()
                    self.gpu_daq.acquire(gpu_photon_slice, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks)
                    gpu_channels = self.gpu_daq.end_acquire()
                    self.gpu_pdf_kernel.accumulate_moments(gpu_channels)

        self.gpu_pdf_kernel.compute_bandwidth(event_channels.hit,
                                              event_channels.t,
                                              event_channels.q,
                                              scale_factor=scale_factor)

    def eval_kernel(self, event_channels,
                    kernel_iterable,
                    trange, qrange,
                    nreps=1, ndaq=1, naverage=1, time_only=True):
        """Returns tuple: 1D array of channel hit counts, 1D array of PDF
        probability densities."""

        self.gpu_pdf_kernel.setup_kernel(event_channels.hit,
                                         event_channels.t,
                                         event_channels.q)
        first_element, kernel_iterable = itertoolset.peek(kernel_iterable)
        if isinstance(first_element, event.Event):
            kernel_iterable = \
                self.photon_generator.generate_events(kernel_iterable)

        # Evaluate likelihood using this bandwidth
        for ev in kernel_iterable:
            gpu_photons = gpu.GPUPhotons(ev.photons_beg, ncopies=nreps)
            gpu_photons.propagate(self.gpu_geometry, self.rng_states,
                                  nthreads_per_block=self.nthreads_per_block,
                                  max_blocks=self.max_blocks)
            for gpu_photon_slice in gpu_photons.iterate_copies():
                for idaq in range(ndaq):
                    self.gpu_daq.begin_acquire()
                    self.gpu_daq.acquire(gpu_photon_slice, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks)
                    gpu_channels = self.gpu_daq.end_acquire()
                    self.gpu_pdf_kernel.accumulate_kernel(gpu_channels)

        return self.gpu_pdf_kernel.get_kernel_eval()

    def __del__(self):
        self.context.pop()


#!/usr/bin/env python

import os
import os.path as osp
import pickle
import shutil

import chainer
import chainercv
import numpy as np
import six
import yaml

from chainer_mask_rcnn.models import MaskRCNNResNet

dataset_dir = chainer.dataset.get_dataset_directory(
    'wkentaro/chainer-mask-rcnn/R-50-C4_1x_caffe2')
dst_file = osp.join(dataset_dir, 'model_final_caffe2_to_chainer.npz')
if osp.exists(dst_file):
    print('Model file already exists: {}'.format(dst_file))
    quit()

src_file = osp.join(dataset_dir, 'model_final.pkl')
if not osp.exists(src_file):
    url = 'https://s3-us-west-2.amazonaws.com/detectron/35858791/12_2017_baselines/e2e_mask_rcnn_R-50-C4_1x.yaml.01_45_57.ZgkA7hPB/output/train/coco_2014_train%3Acoco_2014_valminusminival/generalized_rcnn/model_final.pkl'  # NOQA
    cache_path = chainercv.utils.download.cached_download(url)
    shutil.move(cache_path, src_file)

print('Loading from: {}'.format(src_file))
with open(src_file, 'rb') as f:
    if six.PY2:
        blobs = pickle.load(f)['blobs']
    else:
        blobs = pickle.load(f, encoding='latin-1')['blobs']

model = MaskRCNNResNet(
    n_layers=50,
    n_fg_class=80,
    anchor_scales=[2, 4, 8, 16, 32],
    pretrained_model=None,
    roi_size=14,
)

# /conv1, /bn1
assert all(isinstance(v, np.ndarray) for v in blobs.values())
np.copyto(model.extractor.conv1.W.data, blobs['conv1_w'][:, ::-1])
np.copyto(model.extractor.conv1.b.data, blobs['conv1_b'])
np.copyto(model.extractor.bn1.W.data, blobs['res_conv1_bn_s'])
np.copyto(model.extractor.bn1.b.data, blobs['res_conv1_bn_b'])
# /res2/a
np.copyto(model.extractor.res2.a.conv1.W.data, blobs['res2_0_branch2a_w'])
np.copyto(model.extractor.res2.a.bn1.W.data, blobs['res2_0_branch2a_bn_s'])
np.copyto(model.extractor.res2.a.bn1.b.data, blobs['res2_0_branch2a_bn_b'])
np.copyto(model.extractor.res2.a.conv2.W.data, blobs['res2_0_branch2b_w'])
np.copyto(model.extractor.res2.a.bn2.W.data, blobs['res2_0_branch2b_bn_s'])
np.copyto(model.extractor.res2.a.bn2.b.data, blobs['res2_0_branch2b_bn_b'])
np.copyto(model.extractor.res2.a.conv3.W.data, blobs['res2_0_branch2c_w'])
np.copyto(model.extractor.res2.a.bn3.W.data, blobs['res2_0_branch2c_bn_s'])
np.copyto(model.extractor.res2.a.bn3.b.data, blobs['res2_0_branch2c_bn_b'])
np.copyto(model.extractor.res2.a.conv4.W.data, blobs['res2_0_branch1_w'])
np.copyto(model.extractor.res2.a.bn4.W.data, blobs['res2_0_branch1_bn_s'])
np.copyto(model.extractor.res2.a.bn4.b.data, blobs['res2_0_branch1_bn_b'])
# /res2/b1, /res2/b2
np.copyto(model.extractor.res2.b1.conv1.W.data, blobs['res2_1_branch2a_w'])
np.copyto(model.extractor.res2.b1.bn1.W.data, blobs['res2_1_branch2a_bn_s'])
np.copyto(model.extractor.res2.b1.bn1.b.data, blobs['res2_1_branch2a_bn_b'])
np.copyto(model.extractor.res2.b1.conv2.W.data, blobs['res2_1_branch2b_w'])
np.copyto(model.extractor.res2.b1.bn2.W.data, blobs['res2_1_branch2b_bn_s'])
np.copyto(model.extractor.res2.b1.bn2.b.data, blobs['res2_1_branch2b_bn_b'])
np.copyto(model.extractor.res2.b1.conv3.W.data, blobs['res2_1_branch2c_w'])
np.copyto(model.extractor.res2.b1.bn3.W.data, blobs['res2_1_branch2c_bn_s'])
np.copyto(model.extractor.res2.b1.bn3.b.data, blobs['res2_1_branch2c_bn_b'])
np.copyto(model.extractor.res2.b2.conv1.W.data, blobs['res2_2_branch2a_w'])
np.copyto(model.extractor.res2.b2.bn1.W.data, blobs['res2_2_branch2a_bn_s'])
np.copyto(model.extractor.res2.b2.bn1.b.data, blobs['res2_2_branch2a_bn_b'])
np.copyto(model.extractor.res2.b2.conv2.W.data, blobs['res2_2_branch2b_w'])
np.copyto(model.extractor.res2.b2.bn2.W.data, blobs['res2_2_branch2b_bn_s'])
np.copyto(model.extractor.res2.b2.bn2.b.data, blobs['res2_2_branch2b_bn_b'])
np.copyto(model.extractor.res2.b2.conv3.W.data, blobs['res2_2_branch2c_w'])
np.copyto(model.extractor.res2.b2.bn3.W.data, blobs['res2_2_branch2c_bn_s'])
np.copyto(model.extractor.res2.b2.bn3.b.data, blobs['res2_2_branch2c_bn_b'])
# /res3/a
np.copyto(model.extractor.res3.a.conv1.W.data, blobs['res3_0_branch2a_w'])
np.copyto(model.extractor.res3.a.bn1.W.data, blobs['res3_0_branch2a_bn_s'])
np.copyto(model.extractor.res3.a.bn1.b.data, blobs['res3_0_branch2a_bn_b'])
np.copyto(model.extractor.res3.a.conv2.W.data, blobs['res3_0_branch2b_w'])
np.copyto(model.extractor.res3.a.bn2.W.data, blobs['res3_0_branch2b_bn_s'])
np.copyto(model.extractor.res3.a.bn2.b.data, blobs['res3_0_branch2b_bn_b'])
np.copyto(model.extractor.res3.a.conv3.W.data, blobs['res3_0_branch2c_w'])
np.copyto(model.extractor.res3.a.bn3.W.data, blobs['res3_0_branch2c_bn_s'])
np.copyto(model.extractor.res3.a.bn3.b.data, blobs['res3_0_branch2c_bn_b'])
np.copyto(model.extractor.res3.a.conv4.W.data, blobs['res3_0_branch1_w'])
np.copyto(model.extractor.res3.a.bn4.W.data, blobs['res3_0_branch1_bn_s'])
np.copyto(model.extractor.res3.a.bn4.b.data, blobs['res3_0_branch1_bn_b'])
# /res3/b1, /res3/b2, /res3/b3
np.copyto(model.extractor.res3.b1.conv1.W.data, blobs['res3_1_branch2a_w'])
np.copyto(model.extractor.res3.b1.bn1.W.data, blobs['res3_1_branch2a_bn_s'])
np.copyto(model.extractor.res3.b1.bn1.b.data, blobs['res3_1_branch2a_bn_b'])
np.copyto(model.extractor.res3.b1.conv2.W.data, blobs['res3_1_branch2b_w'])
np.copyto(model.extractor.res3.b1.bn2.W.data, blobs['res3_1_branch2b_bn_s'])
np.copyto(model.extractor.res3.b1.bn2.b.data, blobs['res3_1_branch2b_bn_b'])
np.copyto(model.extractor.res3.b1.conv3.W.data, blobs['res3_1_branch2c_w'])
np.copyto(model.extractor.res3.b1.bn3.W.data, blobs['res3_1_branch2c_bn_s'])
np.copyto(model.extractor.res3.b1.bn3.b.data, blobs['res3_1_branch2c_bn_b'])
np.copyto(model.extractor.res3.b2.conv1.W.data, blobs['res3_2_branch2a_w'])
np.copyto(model.extractor.res3.b2.bn1.W.data, blobs['res3_2_branch2a_bn_s'])
np.copyto(model.extractor.res3.b2.bn1.b.data, blobs['res3_2_branch2a_bn_b'])
np.copyto(model.extractor.res3.b2.conv2.W.data, blobs['res3_2_branch2b_w'])
np.copyto(model.extractor.res3.b2.bn2.W.data, blobs['res3_2_branch2b_bn_s'])
np.copyto(model.extractor.res3.b2.bn2.b.data, blobs['res3_2_branch2b_bn_b'])
np.copyto(model.extractor.res3.b2.conv3.W.data, blobs['res3_2_branch2c_w'])
np.copyto(model.extractor.res3.b2.bn3.W.data, blobs['res3_2_branch2c_bn_s'])
np.copyto(model.extractor.res3.b2.bn3.b.data, blobs['res3_2_branch2c_bn_b'])
np.copyto(model.extractor.res3.b3.conv1.W.data, blobs['res3_3_branch2a_w'])
np.copyto(model.extractor.res3.b3.bn1.W.data, blobs['res3_3_branch2a_bn_s'])
np.copyto(model.extractor.res3.b3.bn1.b.data, blobs['res3_3_branch2a_bn_b'])
np.copyto(model.extractor.res3.b3.conv2.W.data, blobs['res3_3_branch2b_w'])
np.copyto(model.extractor.res3.b3.bn2.W.data, blobs['res3_3_branch2b_bn_s'])
np.copyto(model.extractor.res3.b3.bn2.b.data, blobs['res3_3_branch2b_bn_b'])
np.copyto(model.extractor.res3.b3.conv3.W.data, blobs['res3_3_branch2c_w'])
np.copyto(model.extractor.res3.b3.bn3.W.data, blobs['res3_3_branch2c_bn_s'])
np.copyto(model.extractor.res3.b3.bn3.b.data, blobs['res3_3_branch2c_bn_b'])
# /res4/a
np.copyto(model.extractor.res4.a.conv1.W.data, blobs['res4_0_branch2a_w'])
np.copyto(model.extractor.res4.a.bn1.W.data, blobs['res4_0_branch2a_bn_s'])
np.copyto(model.extractor.res4.a.bn1.b.data, blobs['res4_0_branch2a_bn_b'])
np.copyto(model.extractor.res4.a.conv2.W.data, blobs['res4_0_branch2b_w'])
np.copyto(model.extractor.res4.a.bn2.W.data, blobs['res4_0_branch2b_bn_s'])
np.copyto(model.extractor.res4.a.bn2.b.data, blobs['res4_0_branch2b_bn_b'])
np.copyto(model.extractor.res4.a.conv3.W.data, blobs['res4_0_branch2c_w'])
np.copyto(model.extractor.res4.a.bn3.W.data, blobs['res4_0_branch2c_bn_s'])
np.copyto(model.extractor.res4.a.bn3.b.data, blobs['res4_0_branch2c_bn_b'])
np.copyto(model.extractor.res4.a.conv4.W.data, blobs['res4_0_branch1_w'])
np.copyto(model.extractor.res4.a.bn4.W.data, blobs['res4_0_branch1_bn_s'])
np.copyto(model.extractor.res4.a.bn4.b.data, blobs['res4_0_branch1_bn_b'])
# /res4/b1, /res4/b2, /res4/b3, /res4/b4, /res4/b5
np.copyto(model.extractor.res4.b1.conv1.W.data, blobs['res4_1_branch2a_w'])
np.copyto(model.extractor.res4.b1.bn1.W.data, blobs['res4_1_branch2a_bn_s'])
np.copyto(model.extractor.res4.b1.bn1.b.data, blobs['res4_1_branch2a_bn_b'])
np.copyto(model.extractor.res4.b1.conv2.W.data, blobs['res4_1_branch2b_w'])
np.copyto(model.extractor.res4.b1.bn2.W.data, blobs['res4_1_branch2b_bn_s'])
np.copyto(model.extractor.res4.b1.bn2.b.data, blobs['res4_1_branch2b_bn_b'])
np.copyto(model.extractor.res4.b1.conv3.W.data, blobs['res4_1_branch2c_w'])
np.copyto(model.extractor.res4.b1.bn3.W.data, blobs['res4_1_branch2c_bn_s'])
np.copyto(model.extractor.res4.b1.bn3.b.data, blobs['res4_1_branch2c_bn_b'])
np.copyto(model.extractor.res4.b2.conv1.W.data, blobs['res4_2_branch2a_w'])
np.copyto(model.extractor.res4.b2.bn1.W.data, blobs['res4_2_branch2a_bn_s'])
np.copyto(model.extractor.res4.b2.bn1.b.data, blobs['res4_2_branch2a_bn_b'])
np.copyto(model.extractor.res4.b2.conv2.W.data, blobs['res4_2_branch2b_w'])
np.copyto(model.extractor.res4.b2.bn2.W.data, blobs['res4_2_branch2b_bn_s'])
np.copyto(model.extractor.res4.b2.bn2.b.data, blobs['res4_2_branch2b_bn_b'])
np.copyto(model.extractor.res4.b2.conv3.W.data, blobs['res4_2_branch2c_w'])
np.copyto(model.extractor.res4.b2.bn3.W.data, blobs['res4_2_branch2c_bn_s'])
np.copyto(model.extractor.res4.b2.bn3.b.data, blobs['res4_2_branch2c_bn_b'])
np.copyto(model.extractor.res4.b3.conv1.W.data, blobs['res4_3_branch2a_w'])
np.copyto(model.extractor.res4.b3.bn1.W.data, blobs['res4_3_branch2a_bn_s'])
np.copyto(model.extractor.res4.b3.bn1.b.data, blobs['res4_3_branch2a_bn_b'])
np.copyto(model.extractor.res4.b3.conv2.W.data, blobs['res4_3_branch2b_w'])
np.copyto(model.extractor.res4.b3.bn2.W.data, blobs['res4_3_branch2b_bn_s'])
np.copyto(model.extractor.res4.b3.bn2.b.data, blobs['res4_3_branch2b_bn_b'])
np.copyto(model.extractor.res4.b3.conv3.W.data, blobs['res4_3_branch2c_w'])
np.copyto(model.extractor.res4.b3.bn3.W.data, blobs['res4_3_branch2c_bn_s'])
np.copyto(model.extractor.res4.b3.bn3.b.data, blobs['res4_3_branch2c_bn_b'])
np.copyto(model.extractor.res4.b4.conv1.W.data, blobs['res4_4_branch2a_w'])
np.copyto(model.extractor.res4.b4.bn1.W.data, blobs['res4_4_branch2a_bn_s'])
np.copyto(model.extractor.res4.b4.bn1.b.data, blobs['res4_4_branch2a_bn_b'])
np.copyto(model.extractor.res4.b4.conv2.W.data, blobs['res4_4_branch2b_w'])
np.copyto(model.extractor.res4.b4.bn2.W.data, blobs['res4_4_branch2b_bn_s'])
np.copyto(model.extractor.res4.b4.bn2.b.data, blobs['res4_4_branch2b_bn_b'])
np.copyto(model.extractor.res4.b4.conv3.W.data, blobs['res4_4_branch2c_w'])
np.copyto(model.extractor.res4.b4.bn3.W.data, blobs['res4_4_branch2c_bn_s'])
np.copyto(model.extractor.res4.b4.bn3.b.data, blobs['res4_4_branch2c_bn_b'])
np.copyto(model.extractor.res4.b5.conv1.W.data, blobs['res4_5_branch2a_w'])
np.copyto(model.extractor.res4.b5.bn1.W.data, blobs['res4_5_branch2a_bn_s'])
np.copyto(model.extractor.res4.b5.bn1.b.data, blobs['res4_5_branch2a_bn_b'])
np.copyto(model.extractor.res4.b5.conv2.W.data, blobs['res4_5_branch2b_w'])
np.copyto(model.extractor.res4.b5.bn2.W.data, blobs['res4_5_branch2b_bn_s'])
np.copyto(model.extractor.res4.b5.bn2.b.data, blobs['res4_5_branch2b_bn_b'])
np.copyto(model.extractor.res4.b5.conv3.W.data, blobs['res4_5_branch2c_w'])
np.copyto(model.extractor.res4.b5.bn3.W.data, blobs['res4_5_branch2c_bn_s'])
np.copyto(model.extractor.res4.b5.bn3.b.data, blobs['res4_5_branch2c_bn_b'])
# /rpn: dx, dy, dw, dh -> dy, dx, dh, dw
np.copyto(model.rpn.conv1.W.data, blobs['conv_rpn_w'])
np.copyto(model.rpn.conv1.b.data, blobs['conv_rpn_b'])
W = blobs['rpn_bbox_pred_w']
W = W.reshape(15, 4, 1024, 1, 1)
W = W[:, [1, 0, 3, 2], :, :, :]
W = W.reshape(15 * 4, 1024, 1, 1)
np.copyto(model.rpn.loc.W.data, W)
b = blobs['rpn_bbox_pred_b']
b = b.reshape(15, 4)
b = b[:, [1, 0, 3, 2]]
b = b.reshape(60)
np.copyto(model.rpn.loc.b.data, b)
np.copyto(model.rpn.score.W.data, blobs['rpn_cls_logits_w'])
np.copyto(model.rpn.score.b.data, blobs['rpn_cls_logits_b'])
# /head/res5/a
np.copyto(model.head.res5.a.conv1.W.data, blobs['res5_0_branch2a_w'])
np.copyto(model.head.res5.a.bn1.W.data, blobs['res5_0_branch2a_bn_s'])
np.copyto(model.head.res5.a.bn1.b.data, blobs['res5_0_branch2a_bn_b'])
np.copyto(model.head.res5.a.conv2.W.data, blobs['res5_0_branch2b_w'])
np.copyto(model.head.res5.a.bn2.W.data, blobs['res5_0_branch2b_bn_s'])
np.copyto(model.head.res5.a.bn2.b.data, blobs['res5_0_branch2b_bn_b'])
np.copyto(model.head.res5.a.conv3.W.data, blobs['res5_0_branch2c_w'])
np.copyto(model.head.res5.a.bn3.W.data, blobs['res5_0_branch2c_bn_s'])
np.copyto(model.head.res5.a.bn3.b.data, blobs['res5_0_branch2c_bn_b'])
np.copyto(model.head.res5.a.conv4.W.data, blobs['res5_0_branch1_w'])
np.copyto(model.head.res5.a.bn4.W.data, blobs['res5_0_branch1_bn_s'])
np.copyto(model.head.res5.a.bn4.b.data, blobs['res5_0_branch1_bn_b'])
# /head/res5/b1, /head/res5/b2
np.copyto(model.head.res5.b1.conv1.W.data, blobs['res5_1_branch2a_w'])
np.copyto(model.head.res5.b1.bn1.W.data, blobs['res5_1_branch2a_bn_s'])
np.copyto(model.head.res5.b1.bn1.b.data, blobs['res5_1_branch2a_bn_b'])
np.copyto(model.head.res5.b1.conv2.W.data, blobs['res5_1_branch2b_w'])
np.copyto(model.head.res5.b1.bn2.W.data, blobs['res5_1_branch2b_bn_s'])
np.copyto(model.head.res5.b1.bn2.b.data, blobs['res5_1_branch2b_bn_b'])
np.copyto(model.head.res5.b1.conv3.W.data, blobs['res5_1_branch2c_w'])
np.copyto(model.head.res5.b1.bn3.W.data, blobs['res5_1_branch2c_bn_s'])
np.copyto(model.head.res5.b1.bn3.b.data, blobs['res5_1_branch2c_bn_b'])
np.copyto(model.head.res5.b2.conv1.W.data, blobs['res5_2_branch2a_w'])
np.copyto(model.head.res5.b2.bn1.W.data, blobs['res5_2_branch2a_bn_s'])
np.copyto(model.head.res5.b2.bn1.b.data, blobs['res5_2_branch2a_bn_b'])
np.copyto(model.head.res5.b2.conv2.W.data, blobs['res5_2_branch2b_w'])
np.copyto(model.head.res5.b2.bn2.W.data, blobs['res5_2_branch2b_bn_s'])
np.copyto(model.head.res5.b2.bn2.b.data, blobs['res5_2_branch2b_bn_b'])
np.copyto(model.head.res5.b2.conv3.W.data, blobs['res5_2_branch2c_w'])
np.copyto(model.head.res5.b2.bn3.W.data, blobs['res5_2_branch2c_bn_s'])
np.copyto(model.head.res5.b2.bn3.b.data, blobs['res5_2_branch2c_bn_b'])
# /head/score: dx, dy, dw, dh -> dy, dx, dh, dw
np.copyto(model.head.score.W.data, blobs['cls_score_w'])
np.copyto(model.head.score.b.data, blobs['cls_score_b'])
W = blobs['bbox_pred_w']
W = W.reshape(81, 4, 2048)
W = W[:, [1, 0, 3, 2], :]
W = W.reshape(324, 2048)
# /head/cls_loc
np.copyto(model.head.cls_loc.W.data, W)
b = blobs['bbox_pred_b']
b = b.reshape(81, 4)
b = b[:, [1, 0, 3, 2]]
b = b.reshape(324)
np.copyto(model.head.cls_loc.b.data, b)
# /head/deconv6
np.copyto(model.head.deconv6.W.data, blobs['conv5_mask_w'])
np.copyto(model.head.deconv6.b.data, blobs['conv5_mask_b'])
# /head/mask: remove background class
np.copyto(model.head.mask.W.data, blobs['mask_fcn_logits_w'][1:])
np.copyto(model.head.mask.b.data, blobs['mask_fcn_logits_b'][1:])

# -----------------------------------------------------------------------------

params_src = []
for k, v in sorted(blobs.items()):
    if k.endswith('_momentum'):
        continue
    if k.startswith('fc1000'):
        continue
    if (k.endswith('branch1_b') or k.endswith('branch2a_b') or
            k.endswith('branch2b_b') or k.endswith('branch2c_b')):
        continue
    if k.startswith('mask_fcn_logits_'):
        v = v[1:]
    params_src.extend(v.flatten().tolist())
params_src = np.asarray(params_src)
print(params_src.shape, params_src.min(), params_src.mean(), params_src.max())

params_dst = []
for k, v in model.namedparams():
    v = v.data
    params_dst.extend(v.flatten().tolist())
params_dst = np.asarray(params_dst)
print(params_dst.shape, params_dst.min(), params_dst.mean(), params_dst.max())

# -----------------------------------------------------------------------------

chainer.serializers.save_npz(dst_file, model)
print('Saved to: {}'.format(dst_file))

here = osp.dirname(osp.abspath(__file__))
log_dir = osp.join(here, 'logs/R-50-C4_x1_caffe2_to_chainer')
if not osp.exists(log_dir):
    os.makedirs(log_dir)
link_file = osp.join(log_dir, 'snapshot_model.npz')
if not osp.exists(link_file):
    os.symlink(dst_file, link_file)
yaml_file = osp.join(log_dir, 'params.yaml')
with open(yaml_file, 'w') as f:
    # 0: person ... 79: toothbrush
    with open('coco_class_names.txt') as f2:
        class_names = [n.strip() for n in f2]
    params = dict(
        model='resnet50',
        pooling_func='align',
        roi_size=14,
        mean=(122.7717, 115.9465, 102.9801),
        dataset='coco',
        anchor_scales=(2, 4, 8, 16, 32),
        min_size=800,
        max_size=1333,
        class_names=class_names,
    )
    yaml.safe_dump(params, f, default_flow_style=False)


# -*- coding: utf-8 -*-
"""
dicom2nifti

@author: abrys
"""

from __future__ import print_function
import dicom2nifti.patch_pydicom_encodings
from dicom2nifti.exceptions import ConversionError

dicom2nifti.patch_pydicom_encodings.apply()

import itertools
import os
from math import pow

import logging
import nibabel
import numpy

from pydicom.tag import Tag

from six import string_types

import dicom2nifti.common as common
import dicom2nifti.convert_generic as convert_generic

logger = logging.getLogger(__name__)

def dicom_to_nifti(dicom_input, output_file=None):
    """
    This is the main dicom to nifti conversion fuction for ge images.
    As input ge images are required. It will then determine the type of images and do the correct conversion

    Examples: See unit test

    :param output_file: the filepath to the output nifti file
    :param dicom_input: list with dicom objects
    """
    assert common.is_ge(dicom_input)

    logger.info('Reading and sorting dicom files')
    grouped_dicoms = _get_grouped_dicoms(dicom_input)

    if _is_4d(grouped_dicoms):
        logger.info('Found sequence type: 4D')
        return _4d_to_nifti(grouped_dicoms, output_file)

    logger.info('Assuming anatomical data')
    return convert_generic.dicom_to_nifti(dicom_input, output_file)

def _is_4d(grouped_dicoms):
    """
    Use this function to detect if a dicom series is a ge 4d dataset
    NOTE: Only the first slice will be checked so you can only provide an already sorted dicom directory
    (containing one series)
    """
    # read dicom header
    header = grouped_dicoms[0][0]

    # check if contains multiple stacks
    if len(grouped_dicoms) > 1:
        return True

    return False

def _is_diffusion_imaging(grouped_dicoms):
    """
    Use this function to detect if a dicom series is a ge dti dataset
    NOTE: We already assume this is a 4D dataset
    """
    # we already assume 4D images as input

    # check if contains dti bval information
    bval_tag = Tag(0x0043, 0x1039)  # put this there as this is a slow step and used a lot
    found_bval = False
    for header in list(itertools.chain.from_iterable(grouped_dicoms)):
        if bval_tag in header and int(header[bval_tag].value[0]) != 0:
            found_bval = True
            break
    if not found_bval:
        return False

    return True

def _4d_to_nifti(grouped_dicoms, output_file):
    """
    This function will convert ge 4d series to a nifti
    """

    # Create mosaic block
    logger.info('Creating data block')
    full_block = _get_full_block(grouped_dicoms)

    logger.info('Creating affine')
    # Create the nifti header info
    affine = common.create_affine(grouped_dicoms[0])

    logger.info('Creating nifti')
    # Convert to nifti
    nii_image = nibabel.Nifti1Image(full_block, affine)
    common.set_tr_te(nii_image, float(grouped_dicoms[0][0].RepetitionTime),
                     float(grouped_dicoms[0][0].EchoTime))
    logger.info('Saving nifti to disk %s' % output_file)
    # Save to disk
    if output_file is not None:
        nii_image.to_filename(output_file)

    if _is_diffusion_imaging(grouped_dicoms):
        bval_file = None
        bvec_file = None
        # Create the bval en bevec files
        if output_file is not None:
            base_path = os.path.dirname(output_file)
            base_name = os.path.splitext(os.path.splitext(os.path.basename(output_file))[0])[0]

            logger.info('Creating bval en bvec files')
            bval_file = '%s/%s.bval' % (base_path, base_name)
            bvec_file = '%s/%s.bvec' % (base_path, base_name)
        bval, bvec = _create_bvals_bvecs(grouped_dicoms, bval_file, bvec_file)
        return {'NII_FILE': output_file,
                'BVAL_FILE': bval_file,
                'BVEC_FILE': bvec_file,
                'NII': nii_image,
                'BVAL': bval,
                'BVEC': bvec
                }

    return {'NII_FILE': output_file,
            'NII': nii_image}

def _get_full_block(grouped_dicoms):
    """
    Generate a full datablock containing all timepoints
    """
    # For each slice / mosaic create a data volume block
    data_blocks = []
    for index in range(0, len(grouped_dicoms)):
        logger.info('Creating block %s of %s' % (index + 1, len(grouped_dicoms)))
        data_blocks.append(_timepoint_to_block(grouped_dicoms[index]))

    # Add the data_blocks together to one 4d block
    size_x = numpy.shape(data_blocks[0])[0]
    size_y = numpy.shape(data_blocks[0])[1]
    size_z = numpy.shape(data_blocks[0])[2]
    size_t = len(data_blocks)
    full_block = numpy.zeros((size_x, size_y, size_z, size_t), dtype=data_blocks[0].dtype)
    for index in range(0, size_t):
        if full_block[:, :, :, index].shape != data_blocks[index].shape:
            logger.warning('Missing slices (slice count mismatch between timepoint %s and %s)' % (index - 1, index))
            logger.warning('---------------------------------------------------------')
            logger.warning(full_block[:, :, :, index].shape)
            logger.warning(data_blocks[index].shape)
            logger.warning('---------------------------------------------------------')
            raise ConversionError("MISSING_DICOM_FILES")
        full_block[:, :, :, index] = data_blocks[index]

    return full_block

def _timepoint_to_block(timepoint_dicoms):
    """
    Convert slices to a block of data by reading the headers and appending
    """
    # similar way of getting the block to anatomical however here we are creating the dicom series our selves
    return common.get_volume_pixeldata(timepoint_dicoms)

def _get_grouped_dicoms(dicom_input):
    """
    Search all dicoms in the dicom directory, sort and validate them

    fast_read = True will only read the headers not the data
    """

    # Order all dicom files by InstanceNumber
    dicoms = sorted(dicom_input, key=lambda x: x.InstanceNumber)

    # now group per stack
    grouped_dicoms = [[]]  # list with first element a list
    stack_index = 0

    # loop over all sorted dicoms and sort them by stack
    # for this we use the position and direction of the slices so we can detect a new stack easily
    previous_position = None
    previous_direction = None
    for dicom_ in dicoms:
        current_direction = None
        # if the stack number decreases we moved to the next stack
        if previous_position is not None:
            current_direction = numpy.array(dicom_.ImagePositionPatient) - previous_position
            current_direction = current_direction / numpy.linalg.norm(current_direction)
        if current_direction is not None and \
                previous_direction is not None and \
                not numpy.allclose(current_direction, previous_direction, rtol=0.05, atol=0.05):
            previous_position = numpy.array(dicom_.ImagePositionPatient)
            previous_direction = None
            stack_index += 1
        else:
            previous_position = numpy.array(dicom_.ImagePositionPatient)
            previous_direction = current_direction

        if stack_index >= len(grouped_dicoms):
            grouped_dicoms.append([])
        grouped_dicoms[stack_index].append(dicom_)

    return grouped_dicoms

def _get_bvals_bvecs(grouped_dicoms):
    """
    Write the bvals from the sorted dicom files to a bval file
    """
    # loop over all timepoints and create a list with all bvals and bvecs
    bvals = numpy.zeros([len(grouped_dicoms)], dtype=numpy.int32)
    bvecs = numpy.zeros([len(grouped_dicoms), 3])

    for group_index in range(0, len(grouped_dicoms)):
        dicom_ = grouped_dicoms[group_index][0]
        # 0019:10bb: Diffusion X
        # 0019:10bc: Diffusion Y
        # 0019:10bd: Diffusion Z
        # 0043:1039: B-values (4 values, 1st value is actual B value)

        # bval can be stored both in string as number format in dicom so implement both
        # some workarounds needed for implicit transfer syntax to work
        if isinstance(dicom_[Tag(0x0043, 0x1039)].value, string_types):  # this works for python2.7
            original_bval = float(dicom_[Tag(0x0043, 0x1039)].value.split('\\')[0])
        elif isinstance(dicom_[Tag(0x0043, 0x1039)].value, bytes):  # this works for python3.o
            original_bval = float(dicom_[Tag(0x0043, 0x1039)].value.decode("utf-8").split('\\')[0])
        else:
            original_bval = dicom_[Tag(0x0043, 0x1039)][0]
        original_bvec = numpy.array([0, 0, 0], dtype=numpy.float)
        original_bvec[0] = -float(dicom_[Tag(0x0019, 0x10bb)].value)  # invert based upon mricron output
        original_bvec[1] = float(dicom_[Tag(0x0019, 0x10bc)].value)
        original_bvec[2] = float(dicom_[Tag(0x0019, 0x10bd)].value)

        # Add calculated B Value
        if original_bval != 0:  # only normalize if there is a value
            corrected_bval = original_bval * pow(numpy.linalg.norm(original_bvec), 2)
            if numpy.linalg.norm(original_bvec) != 0:
                normalized_bvec = original_bvec / numpy.linalg.norm(original_bvec)
            else:
                normalized_bvec = original_bvec
        else:
            corrected_bval = original_bval
            normalized_bvec = original_bvec

        bvals[group_index] = int(round(corrected_bval))  # we want the original numbers back as in the protocol
        bvecs[group_index, :] = normalized_bvec

    return bvals, bvecs

def _create_bvals_bvecs(grouped_dicoms, bval_file, bvec_file):
    """
    Write the bvals from the sorted dicom files to a bval file
    """

    # get the bvals and bvecs
    bvals, bvecs = _get_bvals_bvecs(grouped_dicoms)

    # save the found bvecs to the file
    common.write_bval_file(bvals, bval_file)
    common.write_bvec_file(bvecs, bvec_file)

    return bvals, bvecs


import numpy as np
import torch
from torch import nn
import os
from collections import OrderedDict
from torch.autograd import Variable
import torch.nn.functional as F
import EnlightenGAN.util.util as util
from collections import OrderedDict
from torch.autograd import Variable
import itertools
from EnlightenGAN.util.image_pool import ImagePool
from .base_model import BaseModel
import random
from . import networks
from . import fcn
import sys
import bdd.encoding.utils as utils_seg
from utils.perceptual_loss import MSSSIM

msssim_loss = MSSSIM()

def one_hot(index, classes):
    size = index.size()[:1] + (classes,) + index.size()[1:]
    view = index.size()[:1] + (1,) + index.size()[1:]

    mask = torch.Tensor(size).fill_(0).cuda()
    index = index.view(view)
    ones = 1.

    return mask.scatter_(1, index, ones)#[:, 1:, :, :]

class SingleModel(BaseModel):
    def name(self):
        return 'SingleGANModel'

    def initialize(self, opt):
        BaseModel.initialize(self, opt)

        self.multi_D = False
        self.mIoU_delta_mean = 0

        nb = opt.batchSize
        size = opt.fineSize
        self.opt = opt
        self.input_A_origin = self.Tensor(nb, opt.input_nc, size, size)
        self.input_A = self.Tensor(nb, opt.input_nc, size, size)
        self.input_B = self.Tensor(nb, opt.output_nc, size, size)
        self.input_img = self.Tensor(nb, opt.input_nc, size, size)
        self.input_A_gray = self.Tensor(nb, 1, size, size)
        self.mask = torch.cuda.LongTensor(nb, size, size)
        # self.edges_A = torch.cuda.FloatTensor(nb, 1, size, size)
        # self.A_gt = self.Tensor(nb, opt.input_nc, size, size)
        # self.A_boundary = self.Tensor(nb, 1, size, size)

        if opt.vgg > 0:
            self.vgg_loss = networks.PerceptualLoss(opt)
            if self.opt.IN_vgg:
                self.vgg_patch_loss = networks.PerceptualLoss(opt)
                self.vgg_patch_loss.cuda()
            self.vgg_loss.cuda()
            self.vgg = networks.load_vgg16("./model", self.gpu_ids)
            self.vgg.eval()
            for param in self.vgg.parameters():
                param.requires_grad = False
        elif opt.fcn > 0:
            self.fcn_loss = networks.SemanticLoss(opt)
            self.fcn_loss.cuda()
            self.fcn = networks.load_fcn("./model")
            self.fcn.eval()
            for param in self.fcn.parameters():
                param.requires_grad = False
        # load/define networks
        # The naming conversion is different from those used in the paper
        # Code (paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)

        skip = True if opt.skip > 0 else False
        # self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, self.gpu_ids, skip=skip, opt=opt)
        # if model is triain by Segmentation: DataParallel first, then DataParallel
        # if model is trained by EnlightenGAN: filter & load first, then filter & load
        self.netG_A = fcn.FCN(19, 'resnet50', aux=False, se_loss=False, dilated=True, multi_grid=True, multi_dilation=[2, 4, 8])
        self.netG_A = self.netG_A.cuda(device=self.gpu_ids[0])
        # if fix seg part: load pretrained backbone and seg head
        partial = torch.load("/home/chenwy/DynamicDeHaze/cityscapes/fcn_model/res50_di_mg_500px/model_best.pth.tar")['state_dict']
        state = self.netG_A.state_dict()
        pretrained_dict = {k: v for k, v in partial.items() if k in state}
        print("load pretrained segmentation model of length:", len(pretrained_dict))
        state.update(pretrained_dict)
        self.netG_A = torch.nn.DataParallel(self.netG_A)
        # self.netG_A.load_state_dict(state)
        # self.netG_A.module.pretrained.eval()
        # self.netG_A.module.head.eval()

        if self.isTrain:
            use_sigmoid = opt.no_lsgan
            if not self.multi_D:
                self.netD_A = networks.define_D(opt.output_nc, opt.ndf,
                                                opt.which_model_netD,
                                                opt.n_layers_D, opt.norm, use_sigmoid, self.gpu_ids, False)
            else:
                self.netD_As = [networks.define_D(opt.output_nc, opt.ndf,
                                                opt.which_model_netD,
                                                opt.n_layers_D, opt.norm, use_sigmoid, self.gpu_ids, False)
                                                for i in range(5)] # 0: road; 2: building; 8: vegetation; 10:sky; 13: car
            if self.opt.patchD:
                self.netD_P = networks.define_D(opt.input_nc, opt.ndf,
                                            opt.which_model_netD,
                                            opt.n_layers_patchD, opt.norm, use_sigmoid, self.gpu_ids, True)
        if not self.isTrain or opt.continue_train:
            which_epoch = opt.which_epoch
            self.load_network(self.netG_A, 'G_A', which_epoch)
            # self.load_network(self.netG_B, 'G_B', which_epoch)
            if self.isTrain:
                self.load_network(self.netD_A, 'D_A', which_epoch)
                if self.opt.patchD:
                    self.load_network(self.netD_P, 'D_P', which_epoch)

        if self.isTrain:
            self.old_lr = opt.lr
            # self.fake_A_pool = ImagePool(opt.pool_size)
            self.fake_B_pool = ImagePool(opt.pool_size)
            # define loss functions
            if opt.use_wgan:
                self.criterionGAN = networks.DiscLossWGANGP()
            else:
                self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)
            if opt.use_mse:
                self.criterionCycle = torch.nn.MSELoss()
            else:
                self.criterionCycle = torch.nn.L1Loss()
            self.criterionL1 = torch.nn.L1Loss()
            self.criterionIdt = torch.nn.L1Loss()
            # initialize optimizers
            self.optimizer_G = torch.optim.Adam(self.netG_A.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            # fine tune seg & backbone, train header_generator
            # self.optimizer_G = torch.optim.Adam(params=[
            #     {"params": self.netG_A.module.pretrained.parameters(), "lr": opt.lr*0.1, "betas": (opt.beta1, 0.999)},
            #     {"params": self.netG_A.module.head.parameters(), "lr": opt.lr*0.1, "betas": (opt.beta1, 0.999)},
            #     {"params": self.netG_A.module.head_generator.parameters(), "lr": opt.lr, "betas": (opt.beta1, 0.999)},
            # ])
            # fixed seg & backbone, train only header_generator
            # self.optimizer_G = torch.optim.Adam(self.netG_A.module.head_generator.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            if not self.multi_D:
                self.optimizer_D_A = torch.optim.Adam(self.netD_A.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            else:
                self.optimizer_D_A = torch.optim.Adam(
                    list(self.netD_As[0].parameters()) + list(self.netD_As[1].parameters()) + list(self.netD_As[2].parameters()) + list(self.netD_As[3].parameters()) + list(self.netD_As[4].parameters()),
                    lr=opt.lr, betas=(opt.beta1, 0.999))
            if self.opt.patchD:
                self.optimizer_D_P = torch.optim.Adam(self.netD_P.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

        print('---------- Networks initialized -------------')
        networks.print_network(self.netG_A)
        if self.isTrain:
            if not self.multi_D:
                networks.print_network(self.netD_A)
            else:
                networks.print_network(self.netD_As[0])
            if self.opt.patchD:
                networks.print_network(self.netD_P)
        if opt.isTrain:
            self.netG_A.train()
            # self.netG_A.module.head_generator.train()
        else:
            self.netG_A.eval()
        print('-----------------------------------------------')

    def set_input(self, input, domainB=True):
        AtoB = self.opt.which_direction == 'AtoB'
        input_A = input['A' if AtoB else 'B']
        input_img = input['input_img']
        input_A_gray = input['A_gray']

        self.input_A.resize_(input_A.size()).copy_(input_A)
        self.input_A_gray.resize_(input_A_gray.size()).copy_(input_A_gray)
        if domainB:
            input_B = input['B' if AtoB else 'A']
            self.input_B.resize_(input_B.size()).copy_(input_B)
        self.input_img.resize_(input_img.size()).copy_(input_img)
        self.image_paths = input['A_paths' if AtoB else 'B_paths']
        self.mask.resize_(input['mask'].size()).copy_(input['mask'])
        # self.edges_A.resize_(input['edges_A'].size()).copy_(input['edges_A'])
        # self.A_gt.resize_(input['A_gt'].size()).copy_(input['A_gt'])
        # self.A_boundary.resize_(input['A_boundary'].size()).copy_(input['A_boundary'])
        self.category = input['category']

    def set_input_A(self, A, A_gray):#, edges_A):
        self.input_A.resize_(A.size()).copy_(A)
        self.input_A_gray.resize_(A_gray.size()).copy_(A_gray)
        # self.edges_A.resize_(edges_A.size()).copy_(edges_A)

    def set_input_A_origin(self, A):
        self.input_A_origin.resize_(A.size()).copy_(A)

    def test(self):
        self.real_A = Variable(self.input_A, volatile=True)
        self.real_A_gray = Variable(self.input_A_gray, volatile=True)
        if self.opt.noise > 0:
            self.noise = Variable(torch.cuda.FloatTensor(self.real_A.size()).normal_(mean=0, std=self.opt.noise/255.))
            self.real_A = self.real_A + self.noise
        if self.opt.input_linear:
            self.real_A = (self.real_A - torch.min(self.real_A))/(torch.max(self.real_A) - torch.min(self.real_A))
        # print(np.transpose(self.real_A.data[0].cpu().float().numpy(),(1,2,0))[:2][:2][:])
        if self.opt.skip == 1:
            self.fake_B, self.latent_real_A = self.netG_A.forward(self.real_A, self.real_A_gray)
        else:
            self.fake_B = self.netG_A.forward(self.real_A, self.real_A_gray)
        # self.rec_A = self.netG_B.forward(self.fake_B)

        self.real_B = Variable(self.input_B, volatile=True)

    def predict(self, seg=None):
        self.real_A = Variable(self.input_A, volatile=True)
        self.real_A_gray = Variable(self.input_A_gray, volatile=True)
        if self.opt.noise > 0:
            self.noise = Variable(torch.cuda.FloatTensor(self.real_A.size()).normal_(mean=0, std=self.opt.noise/255.))
            self.real_A = self.real_A + self.noise
        if self.opt.input_linear:
            self.real_A = (self.real_A - torch.min(self.real_A))/(torch.max(self.real_A) - torch.min(self.real_A))
        if self.opt.skip == 1:
            self.fake_B, self.latent_real_A, self.seg_real_A = self.netG_A.forward(self.input_A, self.input_A_gray)
            self.fake_B = self.input_A_origin + self.latent_real_A
            self.fake_B = self.fake_B.clamp(-1, 1)
            if seg: self.fake_B_Seg = F.softmax(seg(self.fake_B.clamp(-1, 1))[0], dim=1)
        else:
            self.fake_B = self.netG_A.forward(self.real_A, self.real_A_gray)

        real_A = util.tensor2im(self.real_A.data)
        fake_B = util.tensor2im(self.fake_B.data)
        A_gray = util.atten2im(self.real_A_gray.data)
        return OrderedDict([('real_A', real_A), ('fake_B', fake_B)]), self.fake_B

    # get image paths
    def get_image_paths(self):
        return self.image_paths

    def forward(self, seg, epoch):
        self.real_A = Variable(self.input_A)
        self.real_B = Variable(self.input_B)
        self.real_A_gray = Variable(self.input_A_gray)
        self.real_img = Variable(self.input_img)
        if self.opt.noise > 0:
            self.noise = Variable(torch.cuda.FloatTensor(self.real_A.size()).normal_(mean=0, std=self.opt.noise/255.))
            self.real_A = self.real_A + self.noise
        if self.opt.input_linear:
            self.real_A = (self.real_A - torch.min(self.real_A))/(torch.max(self.real_A) - torch.min(self.real_A))
        if self.opt.skip == 1:
            # self.A_gt_Seg = seg(self.A_gt.clamp(-1, 1))[0]
            self.real_A_Seg = (F.softmax(seg(self.real_A.clamp(-1, 1))[0], dim=1) + F.softmax(seg(self.real_A.clamp(-1, 1).flip(3))[0], dim=1).flip(3)) / 2
            self.fake_B, self.latent_real_A, self.seg_real_A = self.netG_A.forward(self.real_A, self.real_A_gray) # seg_real_A is the seg predicted from G it self (multi-tasking)
            self.fake_B_Seg = F.softmax(seg(self.fake_B.clamp(-1, 1))[0], dim=1) + F.softmax(seg(self.fake_B.clamp(-1, 1).flip(3))[0], dim=1).flip(3)
        else:
            self.fake_B = self.netG_A.forward(self.real_img, self.real_A_gray)
        if self.opt.patchD:
            h = self.real_A.size(2)
            w = self.real_A.size(3)
            self.fake_patch = torch.empty(self.fake_B.size(0), self.fake_B.size(1), self.opt.patchSize, self.opt.patchSize).cuda()
            self.real_patch = torch.empty(self.fake_B.size(0), self.fake_B.size(1), self.opt.patchSize, self.opt.patchSize).cuda()
            self.input_patch = torch.empty(self.fake_B.size(0), self.fake_B.size(1), self.opt.patchSize, self.opt.patchSize).cuda()

            for i in range(self.fake_B.size(0)):
                w_offset = random.randint(0, max(0, w - self.opt.patchSize - 1))
                h_offset = random.randint(0, max(0, h - self.opt.patchSize - 1))
                unique, counts = np.unique(self.mask[i, h_offset:h_offset + self.opt.patchSize, w_offset:w_offset + self.opt.patchSize].detach().cpu().numpy().astype('int32'), return_counts=True)
                self.fake_patch[i] = self.fake_B[i,:, h_offset:h_offset + self.opt.patchSize, w_offset:w_offset + self.opt.patchSize]
                self.real_patch[i] = self.real_B[i,:, h_offset:h_offset + self.opt.patchSize, w_offset:w_offset + self.opt.patchSize]
                self.input_patch[i] = self.real_A[i,:, h_offset:h_offset + self.opt.patchSize, w_offset:w_offset + self.opt.patchSize]

        if self.opt.patchD_3 > 0:
            self.fake_patch_1 = []
            self.real_patch_1 = []
            self.input_patch_1 = []
            w = self.real_A.size(3)
            h = self.real_A.size(2)
            for _ in range(self.opt.patchD_3):
                self.fake_patch_tmp = torch.empty(self.fake_B.size(0), self.fake_B.size(1), self.opt.patchSize, self.opt.patchSize).cuda()
                self.real_patch_tmp = torch.empty(self.fake_B.size(0), self.fake_B.size(1), self.opt.patchSize, self.opt.patchSize).cuda()
                self.input_patch_tmp = torch.empty(self.fake_B.size(0), self.fake_B.size(1), self.opt.patchSize, self.opt.patchSize).cuda()
                for i in range(self.fake_B.size(0)):
                    w_offset = random.randint(0, max(0, w - self.opt.patchSize - 1))
                    h_offset = random.randint(0, max(0, h - self.opt.patchSize - 1))
                    unique, counts = np.unique(self.mask[i, h_offset:h_offset + self.opt.patchSize, w_offset:w_offset + self.opt.patchSize].detach().cpu().numpy().astype('int32'), return_counts=True)
                    self.fake_patch_tmp[i] = self.fake_B[i,:, h_offset:h_offset + self.opt.patchSize, w_offset:w_offset + self.opt.patchSize]
                    self.real_patch_tmp[i] = self.real_B[i,:, h_offset:h_offset + self.opt.patchSize, w_offset:w_offset + self.opt.patchSize]
                    self.input_patch_tmp[i] = self.real_A[i,:, h_offset:h_offset + self.opt.patchSize, w_offset:w_offset + self.opt.patchSize]
                self.fake_patch_1.append(self.fake_patch_tmp)
                self.real_patch_1.append(self.real_patch_tmp)
                self.input_patch_1.append(self.input_patch_tmp)

    def backward_G(self, epoch, seg_criterion=None, A_gt=False):
        # self.loss_G_A = torch.zeros(1).cuda()
        if not self.multi_D:
            pred_fake = self.netD_A.forward(self.fake_B)
            if self.opt.use_wgan:
                self.loss_G_A = -pred_fake.mean()
            elif self.opt.use_ragan:
                pred_real = self.netD_A.forward(self.real_B)
                self.loss_G_A = (self.criterionGAN(pred_real - torch.mean(pred_fake), False) + self.criterionGAN(pred_fake - torch.mean(pred_real), True)) / 2
            else:
                self.loss_G_A = self.criterionGAN(pred_fake, True)
        else:
            self.loss_G_A = 0
            for c in range(5):
                # select by category; if empty: tensor([])
                if (self.category == c).nonzero().size(0) == 0: continue
                pred_fake = self.netD_As[c].forward(torch.index_select(self.fake_B, 0, (self.category == c).nonzero().view(-1).type(torch.cuda.LongTensor)))
                if self.opt.use_wgan:
                    self.loss_G_A += -pred_fake.mean()
                elif self.opt.use_ragan:
                    pred_real = self.netD_As[c].forward(torch.index_select(self.real_B, 0, (self.category == c).nonzero().view(-1).type(torch.cuda.LongTensor)))
                    self.loss_G_A += (self.criterionGAN(pred_real - torch.mean(pred_fake), False) + self.criterionGAN(pred_fake - torch.mean(pred_real), True)) / 2
                else:
                    self.loss_G_A += self.criterionGAN(pred_fake, True)

        loss_G_A = 0
        if self.opt.patchD:
            pred_fake_patch = self.netD_P.forward(self.fake_patch)
            if self.opt.hybrid_loss:
                loss_G_A += self.criterionGAN(pred_fake_patch, True)
            else:
                pred_real_patch = self.netD_P.forward(self.real_patch)
                loss_G_A += (self.criterionGAN(pred_real_patch - torch.mean(pred_fake_patch), False) + self.criterionGAN(pred_fake_patch - torch.mean(pred_real_patch), True)) / 2
        if self.opt.patchD_3 > 0:
            for i in range(self.opt.patchD_3):
                pred_fake_patch_1 = self.netD_P.forward(self.fake_patch_1[i])
                if self.opt.hybrid_loss:
                    loss_G_A += self.criterionGAN(pred_fake_patch_1, True)
                else:
                    pred_real_patch_1 = self.netD_P.forward(self.real_patch_1[i])
                    loss_G_A += (self.criterionGAN(pred_real_patch_1 - torch.mean(pred_fake_patch_1), False) + self.criterionGAN(pred_fake_patch_1 - torch.mean(pred_real_patch_1), True)) / 2

            if not self.opt.D_P_times2:
                self.loss_G_A += loss_G_A/float(self.opt.patchD_3 + 1)
            else:
                self.loss_G_A += loss_G_A/float(self.opt.patchD_3 + 1)*2
        else:
            if not self.opt.D_P_times2:
                self.loss_G_A += loss_G_A
            else:
                self.loss_G_A += loss_G_A*2

        self.loss_G = self.loss_G_A
        if epoch < 0:
            vgg_w = 0
        else:
            if seg_criterion is None: vgg_w = 1
            else: vgg_w = 0
        if vgg_w > 0:
            if self.opt.vgg > 0:
                self.loss_vgg_b = self.vgg_loss.compute_vgg_loss(self.vgg, self.fake_B, self.real_A) * self.opt.vgg if self.opt.vgg > 0 else 0
                if self.opt.patch_vgg:
                    if not self.opt.IN_vgg:
                        loss_vgg_patch = self.vgg_loss.compute_vgg_loss(self.vgg, self.fake_patch, self.input_patch) * self.opt.vgg
                    else:
                        loss_vgg_patch = self.vgg_patch_loss.compute_vgg_loss(self.vgg, self.fake_patch, self.input_patch) * self.opt.vgg
                    if self.opt.patchD_3 > 0:
                        for i in range(self.opt.patchD_3):
                            if not self.opt.IN_vgg:
                                loss_vgg_patch += self.vgg_loss.compute_vgg_loss(self.vgg, self.fake_patch_1[i], self.input_patch_1[i]) * self.opt.vgg
                            else:
                                loss_vgg_patch += self.vgg_patch_loss.compute_vgg_loss(self.vgg, self.fake_patch_1[i], self.input_patch_1[i]) * self.opt.vgg
                        self.loss_vgg_b += loss_vgg_patch/float(self.opt.patchD_3 + 1)
                    else:
                        self.loss_vgg_b += loss_vgg_patch
                self.loss_G = self.loss_G_A + self.loss_vgg_b*vgg_w
            elif self.opt.fcn > 0:
                self.loss_fcn_b = self.fcn_loss.compute_fcn_loss(self.fcn, self.fake_B, self.real_A) * self.opt.fcn if self.opt.fcn > 0 else 0
                if self.opt.patchD:
                    loss_fcn_patch = self.fcn_loss.compute_vgg_loss(self.fcn, self.fake_patch, self.input_patch) * self.opt.fcn
                    if self.opt.patchD_3 > 0:
                        for i in range(self.opt.patchD_3):
                            loss_fcn_patch += self.fcn_loss.compute_vgg_loss(self.fcn, self.fake_patch_1[i], self.input_patch_1[i]) * self.opt.fcn
                        self.loss_fcn_b += loss_fcn_patch/float(self.opt.patchD_3 + 1)
                    else:
                        self.loss_fcn_b += loss_fcn_patch
                self.loss_G = self.loss_G_A + self.loss_fcn_b*vgg_w
            # self.loss_G = self.L1_AB + self.L1_BA

        ## Seg Loss ################################
        if seg_criterion is not None:
            # mIoU of enhanced image
            inter, union = utils_seg.batch_intersection_union(self.fake_B_Seg.data, self.mask, 19)
            idx = union > 0
            IoU = 1.0 * inter[idx] / (np.spacing(1) + union[idx])
            self.mIoU = np.nan_to_num(IoU.mean())

            with torch.no_grad():
                # mIoU of origin image by pretrained Seg Model
                inter, union = utils_seg.batch_intersection_union(self.real_A_Seg.data, self.mask, 19)
                idx = union > 0
                IoU = 1.0 * inter[idx] / (np.spacing(1) + union[idx])
                self.mIoU_ori = np.nan_to_num(IoU.mean())
                self.mIoU_delta_mean = 0.8 * self.mIoU_delta_mean + 0.2 * np.round(self.mIoU-self.mIoU_ori, 3)

                # mIoU of origin image by Generator
                inter, union = utils_seg.batch_intersection_union(self.seg_real_A.data, self.mask, 19)
                idx = union > 0
                IoU = 1.0 * inter[idx] / (np.spacing(1) + union[idx])
                print("mIoU_generator", np.round(np.nan_to_num(IoU.mean()), 3))

            print("G:", self.loss_G.data[0], "mIoU gain:", np.round(self.mIoU-self.mIoU_ori, 3), "mean:", np.round(self.mIoU_delta_mean, 3), "lum:", 255*(1 - self.input_A_gray).mean(), "epoch:", epoch)

            lambd = 3
            self.loss_Seg = seg_criterion(self.fake_B_Seg, self.mask) + lambd * seg_criterion(self.seg_real_A, self.mask)
            self.loss_G += self.loss_Seg
        ############################################
        ## GAN_GT Loss ################################
        if A_gt:
            # msssim = msssim_loss((self.fake_B.clamp(-1, 1)+1)/2*255, (self.A_gt+1)/2*255, weight_map=self.A_boundary)
            l1 = (F.l1_loss((self.fake_B+1)/2*255, (self.A_gt+1)/2*255, reduction='none') * self.A_boundary).mean()
            # self.loss_gt = 3 * msssim + 0.16 * l1
            self.loss_gt = 0.1 * l1
            print("loss_gt", self.loss_gt.data[0])
            self.loss_G += self.loss_gt
        ############################################

        self.loss_G.backward(retain_graph=True)

    def backward_D_basic(self, netD, real, fake, use_ragan):
        # Real
        pred_real = netD.forward(real)
        pred_fake = netD.forward(fake.detach())
        if self.opt.use_wgan:
            loss_D_real = pred_real.mean()
            loss_D_fake = pred_fake.mean()
            loss_D = loss_D_fake - loss_D_real + self.criterionGAN.calc_gradient_penalty(netD, real.data, fake.data)
        elif self.opt.use_ragan and use_ragan:
            loss_D = (self.criterionGAN(pred_real - torch.mean(pred_fake), True) +
                                      self.criterionGAN(pred_fake - torch.mean(pred_real), False)) / 2
        else:
            loss_D_real = self.criterionGAN(pred_real, True)
            loss_D_fake = self.criterionGAN(pred_fake, False)
            loss_D = (loss_D_real + loss_D_fake) * 0.5
        # loss_D.backward()
        return loss_D

    def backward_D_A(self):
        # fake_B = self.fake_B_pool.query(self.fake_B)
        fake_B = self.fake_B
        if not self.multi_D:
            self.loss_D_A = self.backward_D_basic(self.netD_A, self.real_B, fake_B, True)
        else:
            self.loss_D_A = 0
            for c in range(5):
                # select by category; if empty: tensor([])
                if (self.category == c).nonzero().size(0) == 0: continue
                self.loss_D_A += self.backward_D_basic(self.netD_As[c], torch.index_select(self.real_B, 0, (self.category == c).nonzero().view(-1).type(torch.cuda.LongTensor)), torch.index_select(fake_B, 0, (self.category == c).nonzero().view(-1).type(torch.cuda.LongTensor)), True)
        self.loss_D_A.backward()

    def backward_D_P(self):
        if self.opt.hybrid_loss:
            loss_D_P = self.backward_D_basic(self.netD_P, self.real_patch, self.fake_patch, False)
            if self.opt.patchD_3 > 0:
                for i in range(self.opt.patchD_3):
                    loss_D_P += self.backward_D_basic(self.netD_P, self.real_patch_1[i], self.fake_patch_1[i], False)
                self.loss_D_P = loss_D_P/float(self.opt.patchD_3 + 1)
            else:
                self.loss_D_P = loss_D_P
        else:
            loss_D_P = self.backward_D_basic(self.netD_P, self.real_patch, self.fake_patch, True)
            if self.opt.patchD_3 > 0:
                for i in range(self.opt.patchD_3):
                    loss_D_P += self.backward_D_basic(self.netD_P, self.real_patch_1[i], self.fake_patch_1[i], True)
                self.loss_D_P = loss_D_P/float(self.opt.patchD_3 + 1)
            else:
                self.loss_D_P = loss_D_P
        if self.opt.D_P_times2:
            self.loss_D_P = self.loss_D_P*2
        self.loss_D_P.backward()

    def optimize_parameters(self, epoch, step, seg=None, seg_criterion=None):
        late_update = 2
        # forward
        self.forward(seg, epoch)
        # G_A and G_B
        self.backward_G(epoch, seg_criterion=seg_criterion, A_gt=False)
        if step % late_update == 0: # late update
            self.optimizer_G.step()
            self.optimizer_G.zero_grad()
        # D_A
        self.backward_D_A()
        if self.opt.patchD: self.backward_D_P()
        if step % late_update == 0:
            self.optimizer_D_A.step()
            self.optimizer_D_A.zero_grad()
            if self.opt.patchD:
                self.optimizer_D_P.step()
                self.optimizer_D_P.zero_grad()

    def get_current_errors(self, epoch):
        D_A = self.loss_D_A.data[0]
        D_P = self.loss_D_P.data[0] if self.opt.patchD else 0
        G_A = self.loss_G_A.data[0]
        if self.opt.vgg > 0:
            return OrderedDict([('D_A', D_A), ('G_A', G_A), ("D_P", D_P), ("mIoU gain", np.round(self.mIoU-self.mIoU_ori, 3))])
        elif self.opt.fcn > 0:
            fcn = self.loss_fcn_b.data[0]/self.opt.fcn if self.opt.fcn > 0 else 0
            if 'images' in self.opt.name:
                return OrderedDict([('D_A', D_A), ('G_A', G_A), ("fcn", fcn), ("D_P", D_P)])
            else:
                return OrderedDict([('D_A', D_A), ('G_A', G_A), ("fcn", fcn), ("GT", self.loss_gt.data[0]), ("D_P", D_P), ("mIoU gain", np.round(self.mIoU-self.mIoU_ori, 3))])

    def get_current_visuals(self):
        real_A = util.tensor2im(self.real_A.data)
        fake_B = util.tensor2im(self.fake_B.data)
        real_B = util.tensor2im(self.real_B.data)
        if self.opt.skip > 0:
            latent_real_A = util.tensor2im(self.latent_real_A.data)
            latent_show = util.latent2im(self.latent_real_A.data)
            if self.opt.patchD:
                fake_patch = util.tensor2im(self.fake_patch.data)
                real_patch = util.tensor2im(self.real_patch.data)
                if self.opt.patch_vgg:
                    input_patch = util.tensor2im(self.input_patch.data)
                    if not self.opt.self_attention:
                        return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('latent_real_A', latent_real_A),
                                ('latent_show', latent_show), ('real_B', real_B), ('real_patch', real_patch),
                                ('fake_patch', fake_patch), ('input_patch', input_patch)])
                    else:
                        # currently into here
                        self_attention = util.atten2im(self.real_A_gray.data)
                        return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('latent_real_A', latent_real_A),
                                ('latent_show', latent_show), ('real_B', real_B), ('real_patch', real_patch),
                                ('fake_patch', fake_patch), ('input_patch', input_patch), ('self_attention', self_attention)
                                ])
                else:
                    if not self.opt.self_attention:
                        return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('latent_real_A', latent_real_A),
                                ('latent_show', latent_show), ('real_B', real_B), ('real_patch', real_patch),
                                ('fake_patch', fake_patch)])
                    else:
                        self_attention = util.atten2im(self.real_A_gray.data)
                        return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('latent_real_A', latent_real_A),
                                ('latent_show', latent_show), ('real_B', real_B), ('real_patch', real_patch),
                                ('fake_patch', fake_patch), ('self_attention', self_attention)])
            else:
                if not self.opt.self_attention:
                    return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('latent_real_A', latent_real_A),
                                ('latent_show', latent_show), ('real_B', real_B)])
                else:
                    self_attention = util.atten2im(self.real_A_gray.data)
                    return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('real_B', real_B),
                                    ('latent_real_A', latent_real_A), ('latent_show', latent_show),
                                    ('self_attention', self_attention)])
        else:
            if not self.opt.self_attention:
                return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('real_B', real_B)])
            else:
                self_attention = util.atten2im(self.real_A_gray.data)
                return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('real_B', real_B),
                                    ('self_attention', self_attention)])

    def save(self, label):
        self.save_network(self.netG_A, 'G_A', label, self.gpu_ids)
        if not self.multi_D:
            self.save_network(self.netD_A, 'D_A', label, self.gpu_ids)
        else:
            for c in range(5):
                self.save_network(self.netD_As[c], 'D_A'+str(c), label, self.gpu_ids)
        if self.opt.patchD:
            self.save_network(self.netD_P, 'D_P', label, self.gpu_ids)

    def update_learning_rate(self):

        if self.opt.new_lr:
            lr = self.old_lr/2
        else:
            lrd = self.opt.lr / self.opt.niter_decay
            lr = self.old_lr - lrd
        for param_group in self.optimizer_D_A.param_groups:
            param_group['lr'] = lr
        if self.opt.patchD:
            for param_group in self.optimizer_D_P.param_groups:
                param_group['lr'] = lr
        for param_group in self.optimizer_G.param_groups:
            param_group['lr'] = lr

        print('update learning rate: %f -> %f' % (self.old_lr, lr))
        self.old_lr = lr


import os
import copy
import logging
import numpy as np

import torch
import torch.nn.functional as F
from layers import Embedding, masked_softmax, NoisyLinear
from layers import EncoderBlock, CQAttention, AnswerPointer, MergeEmbeddings, AggregationBlock
from generic import to_pt

logger = logging.getLogger(__name__)

class QA_DQN(torch.nn.Module):
    model_name = 'qa_dqn'

    def __init__(self, config, word_vocab, char_vocab, action_space_size):
        super(QA_DQN, self).__init__()
        self.config = config
        self.word_vocab = word_vocab
        self.word_vocab_size = len(word_vocab)
        self.char_vocab = char_vocab
        self.char_vocab_size = len(char_vocab)
        self.action_space_size = action_space_size  # previous, next, stop, ctrl+f <query>
        self.read_config()
        self._def_layers()
        # self.print_parameters()

    def print_parameters(self):
        amount = 0
        for p in self.parameters():
            amount += np.prod(p.size())
        print("total number of parameters: %s" % (amount))
        parameters = filter(lambda p: p.requires_grad, self.parameters())
        amount = 0
        for p in parameters:
            amount += np.prod(p.size())
        print("number of trainable parameters: %s" % (amount))

    def read_config(self):
        # model config
        model_config = self.config['model']
        # word
        self.use_pretrained_embedding = model_config['use_pretrained_embedding']
        self.word_embedding_size = model_config['word_embedding_size']
        self.word_embedding_trainable = model_config['word_embedding_trainable']
        # char
        self.char_embedding_size = model_config['char_embedding_size']
        self.char_embedding_trainable = model_config['char_embedding_trainable']
        self.embedding_dropout = model_config['embedding_dropout']

        self.encoder_layers = model_config['encoder_layers']
        self.encoder_conv_num = model_config['encoder_conv_num']
        self.aggregation_layers = model_config['aggregation_layers']
        self.aggregation_conv_num = model_config['aggregation_conv_num']
        self.block_hidden_dim = model_config['block_hidden_dim']
        self.n_heads = model_config['n_heads']
        self.block_dropout = model_config['block_dropout']
        self.attention_dropout = model_config['attention_dropout']
        self.action_scorer_hidden_dim = model_config['action_scorer_hidden_dim']
        self.action_scorer_softmax = model_config['action_scorer_softmax']
        self.question_answerer_hidden_dim = model_config['question_answerer_hidden_dim']

        self.pretrained_embedding_path = "crawl-300d-2M.vec.h5"
        self.noisy_net = self.config['epsilon_greedy']['noisy_net']

    def _def_layers(self):

        # word embeddings
        if self.use_pretrained_embedding:
            self.word_embedding = Embedding(embedding_size=self.word_embedding_size,
                                            vocab_size=self.word_vocab_size,
                                            id2word=self.word_vocab,
                                            dropout_rate=self.embedding_dropout,
                                            load_pretrained=True,
                                            trainable=self.word_embedding_trainable,
                                            embedding_oov_init="zero",
                                            pretrained_embedding_path=self.pretrained_embedding_path)
        else:
            self.word_embedding = Embedding(embedding_size=self.word_embedding_size,
                                            vocab_size=self.word_vocab_size,
                                            trainable=self.word_embedding_trainable,
                                            dropout_rate=self.embedding_dropout)

        # char embeddings
        self.char_embedding = Embedding(embedding_size=self.char_embedding_size,
                                        vocab_size=self.char_vocab_size,
                                        trainable=self.char_embedding_trainable,
                                        dropout_rate=self.embedding_dropout)

        self.emb_enc = torch.nn.ModuleList([EncoderBlock(conv_num=self.encoder_conv_num, ch_num=self.block_hidden_dim, k=7, block_hidden_dim=self.block_hidden_dim, n_head=self.n_heads, dropout=self.block_dropout) for _ in range(self.encoder_layers)])
        self.merge_embeddings = MergeEmbeddings(block_hidden_dim=self.block_hidden_dim, word_emb_dim=self.word_embedding_size, char_emb_dim=self.char_embedding_size, dropout=self.embedding_dropout)
        self.context_question_attention = CQAttention(block_hidden_dim=self.block_hidden_dim, dropout=self.attention_dropout)
        self.context_question_attention_resizer = torch.nn.Linear(self.block_hidden_dim * 4, self.block_hidden_dim)
        self.model_enc_blks = torch.nn.ModuleList([AggregationBlock(conv_num=self.aggregation_conv_num, ch_num=self.block_hidden_dim,
                                                                    k=5, block_hidden_dim=self.block_hidden_dim,
                                                                    n_head=self.n_heads, dropout=self.block_dropout) for _ in range(self.aggregation_layers)])
        self.answer_pointer = AnswerPointer(block_hidden_dim=self.block_hidden_dim, noisy_net=self.noisy_net)
        encoder_output_dim = self.block_hidden_dim

        linear_function = NoisyLinear if self.noisy_net else torch.nn.Linear

        action_scorer_output_size = 1
        action_scorer_advantage_output_size = self.action_space_size
        action_scorer_ctrlf_output_size = 1
        action_scorer_ctrlf_advantage_output_size = self.word_vocab_size

        if self.action_scorer_hidden_dim > 0:
            self.action_scorer_shared_linear = linear_function(encoder_output_dim, self.action_scorer_hidden_dim)
            action_scorer_input_size = self.action_scorer_hidden_dim
        else:
            action_scorer_input_size = encoder_output_dim

        self.action_scorer_ctrlf = linear_function(action_scorer_input_size, action_scorer_ctrlf_output_size)
        self.action_scorer_ctrlf_advantage = linear_function(action_scorer_input_size, action_scorer_ctrlf_advantage_output_size)

        self.action_scorer_linear = linear_function(action_scorer_input_size, action_scorer_output_size)
        self.action_scorer_linear_advantage = linear_function(action_scorer_input_size, action_scorer_advantage_output_size)

    def get_match_representations(self, doc_encodings, doc_mask, q_encodings, q_mask):
        X = self.context_question_attention(doc_encodings, q_encodings, doc_mask, q_mask)
        M0 = self.context_question_attention_resizer(X)
        M0 = F.dropout(M0, p=self.block_dropout, training=self.training)
        square_mask = torch.bmm(doc_mask.unsqueeze(-1), doc_mask.unsqueeze(1))  # batch x time x time
        for i in range(self.aggregation_layers):
             M0 = self.model_enc_blks[i](M0, doc_mask, square_mask, i * (self.aggregation_conv_num + 2) + 1, self.aggregation_layers)
        return M0

    def representation_generator(self, _input_words, _input_chars):
        embeddings, mask = self.word_embedding(_input_words)  # batch x time x emb
        char_embeddings, _ = self.char_embedding(_input_chars)  # batch x time x nchar x emb
        merged_embeddings = self.merge_embeddings(embeddings, char_embeddings, mask)  # batch x time x emb
        square_mask = torch.bmm(mask.unsqueeze(-1), mask.unsqueeze(1))  # batch x time x time
        for i in range(self.encoder_layers):
            encoding_sequence = self.emb_enc[i](merged_embeddings, square_mask, i * (self.encoder_conv_num + 2) + 1, self.encoder_layers)  # batch x time x enc

        return encoding_sequence, mask

    def action_scorer(self, state_representation_sequence, ctrlf_word_mask):
        # state_representation: batch x time x enc_dim
        # ctrlf_word_mask: batch x vocab
        state_representation, _ = torch.max(state_representation_sequence, 1)
        if self.action_scorer_hidden_dim > 0:
            state_representation = self.action_scorer_shared_linear(state_representation)  # action scorer hidden dim
            state_representation = torch.relu(state_representation)

        a_rank = self.action_scorer_linear(state_representation)  #  batch x 1
        a_rank_advantage = self.action_scorer_linear_advantage(state_representation)  # advantage stream  batch x n_vocab
        a_rank = a_rank + a_rank_advantage - a_rank_advantage.mean(1, keepdim=True)  # combine streams
        if self.action_scorer_softmax:
            a_rank = masked_softmax(a_rank, axis=-1)  # batch x n_vocab

        ctrlf_rank = self.action_scorer_ctrlf(state_representation)  #  batch x 1
        ctrlf_rank_advantage = self.action_scorer_ctrlf_advantage(state_representation)  # advantage stream  batch x n_vocab
        ctrlf_rank_advantage = ctrlf_rank_advantage * ctrlf_word_mask
        ctrlf_rank = ctrlf_rank + ctrlf_rank_advantage - ctrlf_rank_advantage.mean(1, keepdim=True)  # combine streams
        ctrlf_rank = ctrlf_rank * ctrlf_word_mask
        if self.action_scorer_softmax:
            ctrlf_rank = masked_softmax(ctrlf_rank, ctrlf_word_mask, axis=-1)  # batch x n_vocab
        return a_rank, ctrlf_rank

    def answer_question(self, matching_representation_sequence, doc_mask):
        square_mask = torch.bmm(doc_mask.unsqueeze(-1), doc_mask.unsqueeze(1))  # batch x time x time
        M0 = matching_representation_sequence
        M1 = M0
        for i in range(self.aggregation_layers):
             M0 = self.model_enc_blks[i](M0, doc_mask, square_mask, i * (self.aggregation_conv_num + 2) + 1, self.aggregation_layers)
        M2 = M0
        M0 = F.dropout(M0, p=self.block_dropout, training=self.training)
        for i in range(self.aggregation_layers):
             M0 = self.model_enc_blks[i](M0, doc_mask, square_mask, i * (self.aggregation_conv_num + 2) + 1, self.aggregation_layers)
        M3 = M0
        pred = self.answer_pointer(M1, M2, M3, doc_mask)
        return pred

    def reset_noise(self):
        if self.noisy_net:
            self.action_scorer_ctrlf.reset_noise()
            self.action_scorer_ctrlf_advantage.reset_noise()
            self.action_scorer_linear.reset_noise()
            self.action_scorer_linear_advantage.reset_noise()
            self.answer_pointer.zero_noise()

    def zero_noise(self):
        if self.noisy_net:
            self.action_scorer_ctrlf.zero_noise()
            self.action_scorer_ctrlf_advantage.zero_noise()
            self.action_scorer_linear.zero_noise()
            self.action_scorer_linear_advantage.zero_noise()
            self.answer_pointer.zero_noise()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Nov 26 00:20:09 2020

@author: hso
This .py-file follows the structure found in https://debuggercafe.com/faster-rcnn-object-detection-with-pytorch/

## Detect food in images - given pretrained model ##

"""

import torchvision
import numpy
import torch
import argparse
import cv2
import detect_utils
from PIL import Image

# Construct argument parser
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', help='path to input image/video')
parser.add_argument('-m', '--min-size', dest='min_size', default=800,
                    help='minimum input size for the FasterRCNN network')
args = vars(parser.parse_args())

# Load model
#model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True,
                                                    #min_size=args['min_size'])
model = torch.load("model_epoch18.pth",map_location=torch.device('cpu'))

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

image = Image.open(args['input'])
model.eval().to(device)
boxes, classes, labels = detect_utils.predict(image, model, device, 0.8)
image = detect_utils.draw_boxes(boxes, classes, labels, image)
cv2.imshow('Image', image)
save_name = "Pred_img"
cv2.imwrite(f"outputs/{save_name}.jpg", image)
cv2.waitKey(0)