#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 27 15:58:40 2018

@author: lz
"""

from __future__ import absolute_import

#import torch
import torch.nn as nn
import torch.nn.functional as F

import torchvision
from torch.nn import init
#from collections import OrderedDict
#import torch.utils.model_zoo as model_zoo


''' ============================= ResNet =============================== '''
'''
    Modified ResNet model defined in Open-ReID project by @Cysu
    (https://github.com/Cysu/open-reid)
'''
#__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
#           'resnet152']


class ResNet(nn.Module):
    __factory = {
        18: torchvision.models.resnet18,
        34: torchvision.models.resnet34,
        50: torchvision.models.resnet50,
        101: torchvision.models.resnet101,
        152: torchvision.models.resnet152,
    }

    def __init__(self, depth, pretrained=True, cut_at_pooling=False,
                 num_features=128, norm=False, dropout=0, num_classes=0):
        super(ResNet, self).__init__()

        self.depth = depth
        self.pretrained = pretrained
        self.cut_at_pooling = cut_at_pooling

        # Construct base (pretrained) resnet
        if depth not in ResNet.__factory:
            raise KeyError("Unsupported depth:", depth)
        self.base = ResNet.__factory[depth](pretrained=pretrained)

        if not self.cut_at_pooling:
            self.num_features = num_features
            self.norm = norm
            self.dropout = dropout
            self.has_embedding = num_features > 0
            self.num_classes = num_classes

            out_planes = self.base.fc.in_features

            # Append new layers
            if self.has_embedding:
                self.feat = nn.Linear(out_planes, self.num_features)
                self.feat_bn = nn.BatchNorm1d(self.num_features)
                init.kaiming_normal(self.feat.weight, mode='fan_out')
                init.constant(self.feat.bias, 0)
                init.constant(self.feat_bn.weight, 1)
                init.constant(self.feat_bn.bias, 0)
            else:
                # Change the num_features to CNN output channels
                self.num_features = out_planes
            if self.dropout > 0:
                self.drop = nn.Dropout(self.dropout)
            if self.num_classes > 0:
                self.classifier = nn.Linear(self.num_features, self.num_classes)
                init.normal(self.classifier.weight, std=0.001)
                init.constant(self.classifier.bias, 0)

        if not self.pretrained:
            self.reset_params()

    def forward(self, x):
        
        x = self.base.conv1(x)
        x = self.base.bn1(x)
        x = self.base.relu(x)
        
        x = self.base.layer1(x)
        x = self.base.layer2(x)
        x = self.base.layer3(x)
        x = self.base.layer4(x)
        if self.cut_at_pooling:
            return x

        x = F.avg_pool2d(x, x.size()[2:])
        x = x.view(x.size(0), -1)

        if self.has_embedding:
            x = self.feat(x)
            x = self.feat_bn(x)
        if self.norm:
            x /= x.norm(p=2, dim=1, keepdim=True)
        elif self.has_embedding:
            x = F.relu(x)
        if self.dropout > 0:
            x = self.drop(x)
        if self.num_classes > 0:
            x = self.classifier(x)
        return x
    
    def extract(self, x):

        x = self.base.conv1(x)
        x = self.base.bn1(x)
        x = self.base.relu(x)
        
        x = self.base.layer1(x)
        x = self.base.layer2(x)
        x = self.base.layer3(x)
        x = self.base.layer4(x)
        if self.cut_at_pooling:
            return x

        x = F.avg_pool2d(x, x.size()[2:])
        x = x.view(x.size(0), -1)

        if self.has_embedding:
            x = self.feat(x)
            x = self.feat_bn(x)
        if self.norm:
            x /= x.norm(p=2, dim=1, keepdim=True)
        elif self.has_embedding:
            x = F.relu(x)
        return x

    def reset_params(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant(m.weight, 1)
                init.constant(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant(m.bias, 0)

''' ============================= Discriminator =============================== '''
class Discriminator(nn.Module):
    def __init__(self, hash_bit):
        super(Discriminator, self).__init__()
        h_dim = [512, 128]
        self.discriminator = nn.Sequential(
            nn.Linear(hash_bit, h_dim[0]),
            nn.ReLU(),
            nn.Linear(h_dim[0], h_dim[1]),
            nn.ReLU(),
            nn.Linear(h_dim[1], 1),
        )
        self.discriminator.cuda()
        self.reset_params()
    
    def forward(self, x):
        x = self.discriminator(x)
        x = x.mean(0)
        return x.view(1)
    
    def reset_params(self):
        for m in self.discriminator.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant(m.weight, 1)
                init.constant(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant(m.bias, 0)
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 27 15:58:40 2018

@author: lz
"""

from __future__ import absolute_import

#import torch
import torch.nn as nn
import torch.nn.functional as F

import torchvision
from torch.nn import init
#from collections import OrderedDict
#import torch.utils.model_zoo as model_zoo


''' ============================= ResNet =============================== '''
'''
    Modified ResNet model defined in Open-ReID project by @Cysu
    (https://github.com/Cysu/open-reid)
'''
#__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
#           'resnet152']


class ResNet(nn.Module):
    __factory = {
        18: torchvision.models.resnet18,
        34: torchvision.models.resnet34,
        50: torchvision.models.resnet50,
        101: torchvision.models.resnet101,
        152: torchvision.models.resnet152,
    }

    def __init__(self, depth, pretrained=True, cut_at_pooling=False,
                 num_features=128, norm=False, dropout=0, num_classes=0):
        super(ResNet, self).__init__()

        self.depth = depth
        self.pretrained = pretrained
        self.cut_at_pooling = cut_at_pooling

        # Construct base (pretrained) resnet
        if depth not in ResNet.__factory:
            raise KeyError("Unsupported depth:", depth)
        self.base = ResNet.__factory[depth](pretrained=pretrained)

        if not self.cut_at_pooling:
            self.num_features = num_features
            self.norm = norm
            self.dropout = dropout
            self.has_embedding = num_features > 0
            self.num_classes = num_classes

            out_planes = self.base.fc.in_features

            # Append new layers
            if self.has_embedding:
                self.feat = nn.Linear(out_planes, self.num_features)
                self.feat_bn = nn.BatchNorm1d(self.num_features)
                init.kaiming_normal(self.feat.weight, mode='fan_out')
                init.constant(self.feat.bias, 0)
                init.constant(self.feat_bn.weight, 1)
                init.constant(self.feat_bn.bias, 0)
            else:
                # Change the num_features to CNN output channels
                self.num_features = out_planes
            if self.dropout > 0:
                self.drop = nn.Dropout(self.dropout)
            if self.num_classes > 0:
                self.classifier = nn.Linear(self.num_features, self.num_classes)
                init.normal(self.classifier.weight, std=0.001)
                init.constant(self.classifier.bias, 0)

        if not self.pretrained:
            self.reset_params()

    def forward(self, x):
        
        x = self.base.conv1(x)
        x = self.base.bn1(x)
        x = self.base.relu(x)
        
        x = self.base.layer1(x)
        x = self.base.layer2(x)
        x = self.base.layer3(x)
        x = self.base.layer4(x)
        if self.cut_at_pooling:
            return x

        x = F.avg_pool2d(x, x.size()[2:])
        x = x.view(x.size(0), -1)

        if self.has_embedding:
            x = self.feat(x)
            x = self.feat_bn(x)
        if self.norm:
            x /= x.norm(p=2, dim=1, keepdim=True)
        elif self.has_embedding:
            x = F.relu(x)
        if self.dropout > 0:
            x = self.drop(x)
        if self.num_classes > 0:
            x = self.classifier(x)
        return x
    
    def extract(self, x):

        x = self.base.conv1(x)
        x = self.base.bn1(x)
        x = self.base.relu(x)
        
        x = self.base.layer1(x)
        x = self.base.layer2(x)
        x = self.base.layer3(x)
        x = self.base.layer4(x)
        if self.cut_at_pooling:
            return x

        x = F.avg_pool2d(x, x.size()[2:])
        x = x.view(x.size(0), -1)

        if self.has_embedding:
            x = self.feat(x)
            x = self.feat_bn(x)
        if self.norm:
            x /= x.norm(p=2, dim=1, keepdim=True)
        elif self.has_embedding:
            x = F.relu(x)
        return x

    def reset_params(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant(m.weight, 1)
                init.constant(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant(m.bias, 0)

''' ============================= Discriminator =============================== '''
class Discriminator(nn.Module):
    def __init__(self, hash_bit):
        super(Discriminator, self).__init__()
        h_dim = [512, 128]
        self.discriminator = nn.Sequential(
            nn.Linear(hash_bit, h_dim[0]),
            nn.ReLU(),
            nn.Linear(h_dim[0], h_dim[1]),
            nn.ReLU(),
            nn.Linear(h_dim[1], 1),
        )
        self.discriminator.cuda()
        self.reset_params()
    
    def forward(self, x):
        x = self.discriminator(x)
        x = x.mean(0)
        return x.view(1)
    
    def reset_params(self):
        for m in self.discriminator.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant(m.weight, 1)
                init.constant(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant(m.bias, 0)
#!/usr/bin/env python3 # -*- coding: utf-8 -*- Created on Tue Mar 27 15:58:40 2018  @author: lz #import torch #from collections import OrderedDict #import torch.utils.model_zoo as model_zoo ============================= ResNet =============================== Modified ResNet model defined in Open-ReID project by @Cysu     (https://github.com/Cysu/open-reid) #__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', #           'resnet152'] # Construct base (pretrained) resnet # Append new layers # Change the num_features to CNN output channels ============================= Discriminator ===============================