import torch
import tensorflow as tf
import os
import numpy
import torch
import torch.nn as nn
#import kornia as K
from typing import Any, Dict, List, Type, Union
import cv2
import torch.nn.functional as F
import torch.onnx

import tensorflow as tf



# Загрузка весов PyTorch
checkpoint_path = 'D:\PyCharmProjects\Tf_network\conv_net_model.ckpt'
checkpoint = torch.load(checkpoint_path, map_location='cpu')
print(checkpoint.keys())
state_dict = checkpoint


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 свертка без паддинга."""
    return tf.keras.layers.Conv2D(
        filters=out_planes,
        kernel_size=1,
        strides=stride,
        padding='valid',
        use_bias=False)


def conv3x3(in_planes, out_planes, stride=1):
    """3x3 свертка с паддингом."""
    return tf.keras.layers.Conv2D(
        filters=out_planes,
        kernel_size=3,
        strides=stride,
        padding='same',
        use_bias=False)


class BasicBlock(tf.keras.layers.Layer):
    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(in_planes, planes, stride=stride)
        self.bn1 = tf.keras.layers.BatchNormalization()
        self.relu = tf.keras.layers.ReLU()
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = tf.keras.layers.BatchNormalization()

        if stride != 1 or in_planes != planes:
            self.downsample = tf.keras.Sequential([
                conv1x1(in_planes, planes, stride=stride),
                tf.keras.layers.BatchNormalization(),
            ])
        else:
            self.downsample = None

    def call(self, x, training=False):
        identity = x

        y = self.conv1(x)
        y = self.bn1(y, training=training)
        y = self.relu(y)
        y = self.conv2(y)
        y = self.bn2(y, training=training)

        if self.downsample is not None:
            identity = self.downsample(identity, training=training)

        out = self.relu(tf.keras.layers.add([identity, y]))
        return out


class ResNetFPN_8_2(tf.keras.Model):
    """ResNet+FPN, выходные разрешения 1/8 и 1/2."""

    def __init__(self):
        super(ResNetFPN_8_2, self).__init__()

        initial_dim = 128
        block_dims = [128, 196, 256]
        Block = BasicBlock

        self.in_planes = initial_dim

        # Начальный сверточный слой
        self.conv1 = tf.keras.layers.Conv2D(
            filters=initial_dim,
            kernel_size=7,
            strides=2,
            padding='same',
            use_bias=False
        )
        self.bn1 = tf.keras.layers.BatchNormalization()
        self.relu = tf.keras.layers.ReLU()

        # Слои ResNet
        self.layer1 = self._make_layer(Block, block_dims[0], stride=1)
        self.layer2 = self._make_layer(Block, block_dims[1], stride=2)
        self.layer3 = self._make_layer(Block, block_dims[2], stride=2)

        # FPN upsample
        self.layer3_outconv = conv1x1(block_dims[2], block_dims[2])
        self.layer2_outconv = conv1x1(block_dims[1], block_dims[2])
        self.layer2_outconv2 = tf.keras.Sequential([
            conv3x3(block_dims[2], block_dims[2]),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(),
            conv3x3(block_dims[2], block_dims[1]),
        ])
        self.layer1_outconv = conv1x1(block_dims[0], block_dims[1])
        self.layer1_outconv2 = tf.keras.Sequential([
            conv3x3(block_dims[1], block_dims[1]),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(),
            conv3x3(block_dims[1], block_dims[0]),
        ])

        # Финальные слои для преобразования в 1 канал
        self.final_conv_x3 = tf.keras.layers.Conv2D(
            filters=1,
            kernel_size=1,
            padding='same',
            activation='sigmoid'  # Используйте активацию по необходимости
        )
        self.final_conv_x1 = tf.keras.layers.Conv2D(
            filters=1,
            kernel_size=1,
            padding='same',
            activation='sigmoid'
        )

    def _make_layer(self, block, planes, stride=1):
        layers = []
        layers.append(block(self.in_planes, planes, stride=stride))
        self.in_planes = planes
        layers.append(block(self.in_planes, planes, stride=1))
        return tf.keras.Sequential(layers)

    def call(self, x, training=False):
        # ResNet Backbone
        x0 = self.conv1(x)
        x0 = self.bn1(x0, training=training)
        x0 = self.relu(x0)

        x1 = self.layer1(x0, training=training)  # 1/2
        x2 = self.layer2(x1, training=training)  # 1/4
        x3 = self.layer3(x2, training=training)  # 1/8

        # FPN
        x3_out = self.layer3_outconv(x3)

        x2_out = self.layer2_outconv(x2)
        x3_out_2x = tf.image.resize(
            x3_out,
            size=tf.shape(x2_out)[1:3],
            method='bilinear'
        )
        x2_out = x2_out + x3_out_2x
        x2_out = self.layer2_outconv2(x2_out, training=training)

        x1_out = self.layer1_outconv(x1)
        x2_out_2x = tf.image.resize(
            x2_out,
            size=tf.shape(x1_out)[1:3],
            method='bilinear'
        )
        x1_out = x1_out + x2_out_2x
        x1_out = self.layer1_outconv2(x1_out, training=training)

        x3_out = self.final_conv_x3(x3_out)
        x1_out = self.final_conv_x1(x1_out)

        return [x3_out, x1_out]


# Инициализация TensorFlow-модели
tf_model = ResNetFPN_8_2()


# Функция сопоставления
def map_pytorch_to_tf(pytorch_state_dict):
    mapping = {}
    for pytorch_name in pytorch_state_dict:
        tf_name = pytorch_name
        tf_name = tf_name.replace('module.', '')
        tf_name = tf_name.replace('conv', 'conv')
        tf_name = tf_name.replace('bn', 'bn')
        tf_name = tf_name.replace('.', '/')
        tf_name = tf_name.replace('running_mean', 'moving_mean')
        tf_name = tf_name.replace('running_var', 'moving_variance')
        mapping[pytorch_name] = tf_name
    return mapping


mapping = map_pytorch_to_tf(state_dict)

# Перенос весов
for pytorch_name, tf_name in mapping.items():
    pytorch_weight = state_dict[pytorch_name].cpu().numpy()

    # Доступ к слою TensorFlow
    tf_layer = tf_model
    for attr in tf_name.split('/'):
        if hasattr(tf_layer, attr):
            tf_layer = getattr(tf_layer, attr)
        else:
            print(f"Слой '{attr}' не найден в модели TensorFlow.")
            tf_layer = None
            break

    if tf_layer is None:
        continue

    # Перенос весов
    if isinstance(tf_layer, tf.keras.layers.Conv2D):
        if 'weight' in pytorch_name:
            pytorch_weight = np.transpose(pytorch_weight, (2, 3, 1, 0))
            tf_layer.kernel.assign(pytorch_weight)
        elif 'bias' in pytorch_name:
            tf_layer.bias.assign(pytorch_weight)
    elif isinstance(tf_layer, tf.keras.layers.BatchNormalization):
        if 'weight' in pytorch_name:
            tf_layer.gamma.assign(pytorch_weight)
        elif 'bias' in pytorch_name:
            tf_layer.beta.assign(pytorch_weight)
        elif 'moving_mean' in pytorch_name:
            tf_layer.moving_mean.assign(pytorch_weight)
        elif 'moving_variance' in pytorch_name:
            tf_layer.moving_variance.assign(pytorch_weight)
    else:
        print(f"Тип слоя {type(tf_layer)} не обрабатывается.")

# Сохранение модели с загруженными весами
tf_model.save_weights('resnet_fpn.weights.h5')




odict_keys(['conv1.weight', 'bn1.weight', 'bn1.bias', 'bn1.running_mean', 'bn1.running_var', 'bn1.num_batches_tracked', 'layer1.0.conv1.weight', 'layer1.0.conv2.weight', 'layer1.0.bn1.weight', 'layer1.0.bn1.bias', 'layer1.0.bn1.running_mean', 'layer1.0.bn1.running_var', 'layer1.0.bn1.num_batches_tracked', 'layer1.0.bn2.weight', 'layer1.0.bn2.bias', 'layer1.0.bn2.running_mean', 'layer1.0.bn2.running_var', 'layer1.0.bn2.num_batches_tracked', 'layer1.1.conv1.weight', 'layer1.1.conv2.weight', 'layer1.1.bn1.weight', 'layer1.1.bn1.bias', 'layer1.1.bn1.running_mean', 'layer1.1.bn1.running_var', 'layer1.1.bn1.num_batches_tracked', 'layer1.1.bn2.weight', 'layer1.1.bn2.bias', 'layer1.1.bn2.running_mean', 'layer1.1.bn2.running_var', 'layer1.1.bn2.num_batches_tracked', 'layer2.0.conv1.weight', 'layer2.0.conv2.weight', 'layer2.0.bn1.weight', 'layer2.0.bn1.bias', 'layer2.0.bn1.running_mean', 'layer2.0.bn1.running_var', 'layer2.0.bn1.num_batches_tracked', 'layer2.0.bn2.weight', 'layer2.0.bn2.bias', 'layer2.0.bn2.running_mean', 'layer2.0.bn2.running_var', 'layer2.0.bn2.num_batches_tracked', 'layer2.0.downsample.0.weight', 'layer2.0.downsample.1.weight', 'layer2.0.downsample.1.bias', 'layer2.0.downsample.1.running_mean', 'layer2.0.downsample.1.running_var', 'layer2.0.downsample.1.num_batches_tracked', 'layer2.1.conv1.weight', 'layer2.1.conv2.weight', 'layer2.1.bn1.weight', 'layer2.1.bn1.bias', 'layer2.1.bn1.running_mean', 'layer2.1.bn1.running_var', 'layer2.1.bn1.num_batches_tracked', 'layer2.1.bn2.weight', 'layer2.1.bn2.bias', 'layer2.1.bn2.running_mean', 'layer2.1.bn2.running_var', 'layer2.1.bn2.num_batches_tracked', 'layer3.0.conv1.weight', 'layer3.0.conv2.weight', 'layer3.0.bn1.weight', 'layer3.0.bn1.bias', 'layer3.0.bn1.running_mean', 'layer3.0.bn1.running_var', 'layer3.0.bn1.num_batches_tracked', 'layer3.0.bn2.weight', 'layer3.0.bn2.bias', 'layer3.0.bn2.running_mean', 'layer3.0.bn2.running_var', 'layer3.0.bn2.num_batches_tracked', 'layer3.0.downsample.0.weight', 'layer3.0.downsample.1.weight', 'layer3.0.downsample.1.bias', 'layer3.0.downsample.1.running_mean', 'layer3.0.downsample.1.running_var', 'layer3.0.downsample.1.num_batches_tracked', 'layer3.1.conv1.weight', 'layer3.1.conv2.weight', 'layer3.1.bn1.weight', 'layer3.1.bn1.bias', 'layer3.1.bn1.running_mean', 'layer3.1.bn1.running_var', 'layer3.1.bn1.num_batches_tracked', 'layer3.1.bn2.weight', 'layer3.1.bn2.bias', 'layer3.1.bn2.running_mean', 'layer3.1.bn2.running_var', 'layer3.1.bn2.num_batches_tracked', 'layer3_outconv.weight', 'layer2_outconv.weight', 'layer2_outconv2.0.weight', 'layer2_outconv2.1.weight', 'layer2_outconv2.1.bias', 'layer2_outconv2.1.running_mean', 'layer2_outconv2.1.running_var', 'layer2_outconv2.1.num_batches_tracked', 'layer2_outconv2.3.weight', 'layer1_outconv.weight', 'layer1_outconv2.0.weight', 'layer1_outconv2.1.weight', 'layer1_outconv2.1.bias', 'layer1_outconv2.1.running_mean', 'layer1_outconv2.1.running_var', 'layer1_outconv2.1.num_batches_tracked', 'layer1_outconv2.3.weight'])
Слой 'weight' не найден в модели TensorFlow.
Слой 'weight' не найден в модели TensorFlow.
Слой 'bias' не найден в модели TensorFlow.
Слой 'num_batches_tracked' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой 'weight' не найден в модели TensorFlow.
Слой 'weight' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '3' не найден в модели TensorFlow.
Слой 'weight' не найден в модели TensorFlow.
Слой '0' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '1' не найден в модели TensorFlow.
Слой '3' не найден в модели TensorFlow.


я хочу перенести веса из pytorch в tensorflow, помоги