Несоответствие между формой входных и выходных данных Unet.

Я работаю над моделью U-Net с двумя выходами: сегментация и классификация. Мой набор данных имеет 110 классов объектов в формате coco, и я использую TensorFlow/Keras для обучения модели. Сегментационная часть модели выдает классификации на уровне пикселей, а голова классификации предназначена для предсказания класса объекта. Однако, когда я пытаюсь обучить модель, я сталкиваюсь со следующей ошибкой:

import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
import cv2
import json
import os
from pycocotools.coco import COCO

# Убедитесь, что пути к данным COCO верны
train_images_path = "/content/drive/MyDrive/fyp/FoodInsSeg/images/train"
train_annotations_path = "/content/drive/MyDrive/fyp/FoodInsSeg/annotations/Train.json"
# val_images_path = "/path/to/val/images"
# val_annotations_path = "/path/to/val/annotations.json"

import tensorflow as tf
import numpy as np
import cv2
from pycocotools.coco import COCO
from tensorflow.keras.utils import to_categorical

# Загрузка аннотаций набора данных COCO
train_coco = COCO(train_annotations_path)
class CustomDataset(tf.keras.utils.Sequence):
    def __init__(self, image_paths, annotations, coco, num_categories=110, batch_size=1, shuffle=True):
        self.image_paths = image_paths
        self.annotations = annotations
        self.coco = coco
        self.num_categories = num_categories  # Количество категорий для классификации
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.on_epoch_end()

    def on_epoch_end(self):
        self.indices = np.arange(len(self.image_paths))
        if self.shuffle:
            np.random.shuffle(self.indices)

    def __len__(self):
        return len(self.image_paths) // self.batch_size

    def __getitem__(self, index):
        batch_indices = self.indices[index * self.batch_size:(index + 1) * self.batch_size]
        batch_image_paths = [self.image_paths[i] for i in batch_indices]
        batch_annotations = [self.annotations[i] for i in batch_indices]

        images = []
        masks = []
        labels = []

        # Обработка каждого изображения и аннотации независимо
        for img_path, ann_id in zip(batch_image_paths, batch_annotations):
            # Загрузка изображения
            img = tf.image.decode_jpeg(tf.io.read_file(img_path), channels=3)
            img = tf.image.convert_image_dtype(img, tf.float32)

            # Загрузка маски без изменения размера
            ann = self.coco.loadAnns([ann_id])[0]
            mask = self.coco.annToMask(ann)
            category_id = ann['category_id']

            # Расширение маски для добавления размерности канала
            mask = np.expand_dims(mask, axis=-1)
            mask = mask * category_id  # Умножение на id категории (для многоклассовой классификации)

            # One-hot кодирование меток классификации
            one_hot_label = to_categorical(category_id, num_classes=self.num_categories)

            # Добавление в списки батча
            images.append(img)
            masks.append(mask)
            labels.append(one_hot_label)

        # Возвращаем изображение и маску как один батч
        return np.array(images), (np.array(masks), np.array(labels))

# Обернуть набор данных в tf.data.Dataset для совместимости с model.fit
def coco_dataset_loader(annotation_file, img_dir, batch_size):
    dataset = CustomDataset(img_dir, annotation_file, train_coco, batch_size=batch_size)
    return tf.data.Dataset.from_generator(
        lambda: dataset,
        output_signature=(
            tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.float32),   # Динамические размеры изображений
            (
                tf.TensorSpec(shape=(None, None, None, 1), dtype=tf.int32),  # Маски сегментации
                tf.TensorSpec(shape=(None,), dtype=tf.int32)                # Метки классификации
            )
        )
    )

batch_size = 4

# Инициализация набора данных без изменения размера или дополнения
train_dataset = coco_dataset_loader(train_annotations_path, train_images_path, batch_size=batch_size) # val_dataset = CustomCocoDataset(val_annotations_path, val_images_path)

def unet_model(input_size, num_classes, num_categories):
    inputs = layers.Input(input_size)

    # Кодировщик (уменьшение размерности)
    conv1 = layers.Conv2D(64, 3, activation='relu', padding='same')(inputs)
    conv1 = layers.Conv2D(64, 3, activation='relu', padding='same')(conv1)
    pool1 = layers.MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = layers.Conv2D(128, 3, activation='relu', padding='same')(pool1)
    conv2 = layers.Conv2D(128, 3, activation='relu', padding='same')(conv2)
    pool2 = layers.MaxPooling2D(pool_size=(2, 2))(conv2)

    # Узкое место
    conv3 = layers.Conv2D(256, 3, activation='relu', padding='same')(pool2)

    # Декодер (увеличение размерности)
    up2 = layers.Conv2DTranspose(128, 2, strides=(2, 2), padding='same')(conv3)
    merge2 = layers.concatenate([conv2, up2])
    conv4 = layers.Conv2D(128, 3, activation='relu', padding='same')(merge2)

    up1 = layers.Conv2DTranspose(64, 2, strides=(2, 2), padding='same')(conv4)
    merge1 = layers.concatenate([conv1, up1])
    conv5 = layers.Conv2D(64, 3, activation='relu', padding='same')(merge1)

    # Выходной слой для сегментации
    seg_output = layers.Conv2D(num_classes, 1, activation='softmax', name="segmentation")(conv5)

    # Голова классификации
    flat = layers.GlobalAveragePooling2D()(conv3)
    class_output = layers.Dense(num_categories, activation='softmax', name="classification")(flat)

    model = models.Model(inputs, [seg_output, class_output])
    return model
input_shape = (None, None, 3)  # Позволить динамические размеры входа
# Создание модели
model = unet_model(input_size=input_shape, num_classes=2, num_categories=110)

# Компилирование модели
model.compile(optimizer="adam",
              loss={'segmentation': 'sparse_categorical_crossentropy', 'classification': 'categorical_crossentropy'},
              metrics={'segmentation': 'accuracy', 'classification': 'accuracy'})

# Обучение модели на первом наборе данных
model.fit(train_dataset, validation_data=train_dataset, epochs=10)

# Сохранение модели после обучения на первом наборе данных
model.save('/content/drive/MyDrive/fyp/unet/model_unet.h5')

# Загрузка модели и дообучение на втором наборе данных
# model = tf.keras.models.load_model('/path/to/save/model_unet.h5')
# train_dataset_2 = CustomCocoDataset('/path/to/second/dataset/annotations.json', '/path/to/second/dataset/images')
# val_dataset_2 = CustomCocoDataset('/path/to/second/val/annotations.json', '/path/to/second/val/images')

# model.fit(train_dataset_2, validation_data=val_dataset_2, epochs=10)

# # Сохранение финальной модели
# model.save('/path/to/save/final_model_unet.h5')