Прежде чем мы сможем начать обучение модели сегментации изображения, нам нужны хорошие обучающие данные. В частности, нам нужны хорошие размеченные обучающие данные. Маркировка изображений для сегментации изображений требует создания масок; цветные попиксельные области, которые отличают один объект от другого. Создание таких масок — очень трудоемкое предприятие. Поэтому вместо того, чтобы создавать базу данных самостоятельно, мы искали существующий набор данных. Я использовал набор данных 3D-FUTURE: 3D Furniture shape with TextURE, созданный Fu et al., 2021, доступный по запросу в Alibaba Cloud. Этот набор данных использовался для семинара Alibaba 3D Artificial Challenge в 2020 году и состоял из 20 240 изображений виртуальных помещений, примерно 10 000 из которых содержат диван. Несмотря на то, что 3D-рендеринг изображений был далеко за пределами моей цели и возможностей, я мог использовать набор данных, поскольку все изображения были помечены для сегментации экземпляра, что означает, что каждый предмет мебели (экземпляр) получил свою собственную цветовую метку и маску, используя COCO-формат.

Когда мы получили набор данных, он имел следующую структуру каталогов:

3D-FUTURE-scene.zip
├── train
│    ├── image
|    │    ├── 0000000.jpg
|    │    ├── 0000001.jpg
│    │    ├── ...
│    │
│    └── idmap
|         ├── 0000000.png
|         ├── 0000001.png
│         ├── ...
│ 
├── test
│    ├── image
|    │    ├── 0000000.jpg
|    │    ├── 0000001.jpg
│    │    ├── ...
│    │    
│    └── idmap
|         ├── 0000000.png
|         ├── 0000001.png
│         ├── ...
│
└── GT
     ├── model_infos.json
     ├── train_set.json
     └── test_set.json

Каждая из папок train и test содержит две подпапки: папку изображений, содержащую все исходные изображения, и папку idmap, содержащую изображения с цветовой кодировкой, которые сегментируют все объекты мебели на изображениях. Осторожно, это изображение отличается от файлов аннотаций, что приведет к проблемам, которые мы обсудим позже. Кроме того, папка Ground Truth (GT) содержит все фактические файлы аннотаций, которые мы будем использовать. Первый файл в этой папке, model_info.json, описывает стиль, тему, материал и категорию каждой мебели, кроме светильников. Мы не будем использовать этот файл, так как стиль оригинальной мебели не представляет интереса для этого проекта. Затем файлы train_set.json и test_set.json основаны на формате COCO и описывают категории, изображения и аннотации. Они соответствуют следующему формату:

train_set.json / test_set.json:
{
  categories: [
    {
      id: int,
      category_name: str,
      fine-grained category name: str,
    },
    ...
  ]
  images: [
    {
      id: int,
      width: int,
      height: int,
      file_name: str,
    },
    ...
  ],
  annotations: [
    {
      id: int,
      image_id: int,
      category_id: int,
      segmentation: RLE or [polygon],
      area: float,
      bbox: [x,y,width,height],
      model_id: str,
      texture_id: str,
      pose: list
      fov: float,
      style: str,
      theme: str,
      material: str
    },
    ...
  ],
}

Здесь основное внимание уделяется аннотациям в этих файлах. Их можно рассматривать как таблицу, в которой каждая запись представляет собой предмет мебели. Для этого объекта мы получаем идентификатор изображения, на котором оно появляется, идентификатор категории объекта (кресло, диван, кровать размера king-size и т. д.), многоугольник сегментации и много дополнительной информации, которую мы не будем использовать. Теперь, получив доступ к аннотациям данных обучения и тестирования с помощью полезного пакета Pycocotools, мы можем выбрать все изображения в наборе данных, которые содержат диван, и визуализировать их бинарные маски, как показано в коде ниже:

import numpy as np
from pycocotools.coco import COCO
from pycocotools import mask as coco_mask
from PIL import Image

# Read the annotation file
coco = COCO('train_set.json')
# Get all subcategories that fall under the category 'Sofa' 
cat_ids = coco.getCatIds(catNms = 'Sofa')
# Get all image ids that contain a sofa
img_ids = [coco.getImgIds(catIds = cat_id) for cat_id in cat_ids]

### Let us now visualise the binary mask for the first image
img_id = img_ids[0]
# Obtain all of the annotations ids in the image
ann_ids = coco.getAnnIds(imgIds=img_id)
# Get the objects from these annotation ids
targets = coco.loadAnns(ann_ids)
# Get the file name and measurements
file_name = coco.loadImgs(img_id)[0]["file_name"]
img_w, img_h = (
    coco.loadImgs(img_id)[0]["width"],
    coco.loadImgs(img_id)[0]["height"],
    )
# Select only the masks from the sofas
masks = np.zeros((img_w, img_h), dtype=np.uint8)
for target in targets:
    if target["category_id"] in cat_ids:
    	# Get object polygons and convert them to a mask
        polygons = target["segmentation"]  
        rles = coco_mask.frPyObjects(polygons, img_h, img_w)
        mask = coco_mask.decode(rles)
        if len(mask.shape) < 3:
            mask = mask[..., None]
        mask = mask.any(axis=2)
        # Add all masks together into a binary mask
        masks[mask] = 1

# Convert the mask to an image
mask = Image.fromarray(masks)

Этот код генерирует маску ниже:

Исходное изображение (слева) и сгенерированная маска дивана (справа)

Теперь, конечно, не может быть радуги без дождя. Итак, к сожалению, мы заметили некоторые проблемы в данных во время изучения и обучения.

Прежде всего, мы увидели, что предоставленные файлы аннотаций не всегда были полными. Кажется, существует несоответствие между сегментацией экземпляра, визуализируемой на карте идентификаторов каждого изображения, и фактическими масками, представленными в файлах аннотаций. Эта проблема возникает примерно на 10% изображений.

Несмотря на то, что карты ID казались правильными, реконструкция файлов аннотаций в стиле COCO была невозможна, поскольку цветовое кодирование казалось несовместимым между изображениями (например, диван мог быть окрашен в зеленый оттенок на одном изображении, а другой диван мог быть замаскирован). используя пурпурный оттенок на другой).
Как бы мы ни старались избежать этого, некоторая ручная обработка обучающих данных все же была необходима. После тщательной обрезки случайного выбора изображений мы создали чистый набор данных примерно из 2000 изображений для обучения и проверки.

Вскоре обнаружилась вторая проблема. Различие между диванами и стульями, изначально двумя разными категориями в наборе данных сегментации экземпляров, было очень расплывчатым и часто непоследовательным.

Например, взгляните на изображения ниже. Можете ли вы различить их?

Решение

Первые три помечены как стулья, а последние три помечены как кушетки.

Тем не менее, поскольку мой собственный мозг все еще пытается отличить одно от другого, а модель всегда будет смещена в сторону человеческой маркировки, мы допустили это небольшое несовершенство.

Что такое «диван»?

Теперь, когда у нас есть необходимые инструменты, давайте создадим нашего помощника по дизайну интерьера! Однако не слишком увлекайтесь и не начинайте программировать во имя скорости. Несколько поисков в Google показывают нам, что несколько человек столкнулись и решили эту проблему до нас. И, как, наверное, сделал бы каждый программист, я прислушался к ленивому голосу в затылке: «Зачем создавать что-то новое с нуля, если кто-то уже сделал это лучше?».

Этот изящный метод называется трансферным обучением, и благодаря обширному пакету segmentation-models Павла Якубовского мы можем вооружить нашу модель предварительно обученными знаниями из ImageNet. Этот пакет PyPI предлагает широкий выбор архитектур моделей и магистральных сетей. Мы решили сосредоточиться на архитектуре U-net, известной своим успехом в задачах биомедицинского компьютерного зрения, за счет использования структуры кодер-декодер с пропуском соединения между ними. В части кодировщика изображение будет закодировано как вектор признаков. Мы заимствовали это представление из общих моделей классификации, которые упоминаются как основы. В качестве основы мы экспериментировали с несколькими различными вариантами, такими как Resnet, Efficientnet, Mobilnet и Densenet. Затем часть декодера будет повышать дискретизацию вектора признаков обратно в изображение. Пропуски соединений между этими двумя частями помогают уточнить маски сегментации.

При построении нашей модели мы также должны определить оптимизатор, функцию активации, функцию потерь и метрики оценки. При настройке нашей модели я уделил особое внимание расширяемости кода, чтобы он был достаточно универсальным, чтобы его можно было применять для сегментации большего количества классов. В результате получился код, как показано ниже.

В первую очередь определяется количество классов и увеличивается на единицу в случае мультиклассовой сегментации, так как фон добавляется как дополнительный класс. В качестве функции активации мы используем сигмовидную активацию в бинарном случае и расширяем ее до функции softmax в мультиклассовом случае. Кроме того, мы выбрали популярный оптимизатор Adam, потому что он сочетает в себе преимущества алгоритма адаптивного градиента и среднеквадратического распространения. В качестве функции проигрыша мы используем комбинацию проигрыша в кости и фокусного проигрыша. Первый обрабатывает дисбаланс данных и позволяет нам устанавливать собственные веса для разных классов. Это не было необходимо в нашем приложении, но было оставлено как вариант для дальнейшего расширения на большее количество классов. Позднее помогает сосредоточиться на классах, которые труднее предсказать. В случае бинарной сегментации мы используем бинарную фокальную потерю, которая обобщает бинарную кросс-энтропию за счет введения гиперпараметра, называемого параметром фокусировки, который позволяет более строго наказывать трудно поддающиеся классификации примеры по сравнению с легко классифицируемые примеры. Взвешивание этих двух функций потерь также может быть установлено по желанию пользователя. Наконец, в качестве показателей точности мы используем процент пересечения над объединением и оценку F1. Оценка IoU интересна, потому что она очень интуитивно понятна, а оценка F1 важна, потому что на нее не влияет дисбаланс классов.

import segmentation_models as sm

# Define our segmentation model
classes = ['Sofa']
n_classes = 1 if len(classes) == 1 else (len(classes) + 1)

# Define the activation function
activation = "sigmoid" if n_classes == 1 else "softmax"

# Pick a backbone
backbone = 'Resnet50'

# Create the Unet model
model = sm.Unet(backbone, classes=n_classes, activation=activation)

# Define the optomizer
optim = keras.optimizers.Adam(LR)

# Define a loss function
dice_loss = sm.losses.DiceLoss()
focal_loss = (sm.losses.BinaryFocalLoss() if n_classes == 1
				else sm.losses.CategoricalFocalLoss())
total_loss = dice_loss + (1 * focal_loss)

metrics = [sm.metrics.IOUScore(threshold=0.5),
		   sm.metrics.FScore(threshold=0.5)]

# Compile the model with the defined optimizer, loss and metrics
model.compile(optim, total_loss, metrics)

Прежде чем мы начнем обучение, мы должны знать, что эти модели имеют огромное пространство параметров, которое легко затмевает наш небольшой набор данных (особенно если мы тренируемся только на правильно помеченном подмножестве, которое мы создали). Следовательно, мы используем методы увеличения данных, включая кадрирование, добавление шума, масштабирование, смещение, вращение, зеркальное отображение и многое другое. Это позволяет нам синтетически увеличить наш набор данных. Кроме того, мы также внедрили планировщик экспоненциально уменьшающейся скорости обучения и критерий ранней остановки (по точности проверки), чтобы избежать переобучения.

Затем комбинация этих методов приводит к модели Resnet50, обученной в течение 20 эпох в нашем небольшом подмножестве, с проверочной оценкой F1 88%. На представленной ниже визуализации (изображений в проверочном наборе) мы видим, что модель ведет себя довольно хорошо, но все еще борется со стульями и подушками. Хотя модель, вероятно, можно улучшить еще больше за счет дальнейшей очистки данных и более длительного обучения, это показалось хорошей отправной точкой для нашего помощника по искусственному интеллекту.

Как поменять диван?

Поскольку ИИ-помощник теперь (более или менее) умеет определять диваны на изображении, нам нужно научить его менять стиль. Опять же, мы будем использовать уже существующие модели, а не обучать их самостоятельно.

  • Первая модель, которую мы применили к нашим диванам, — это модель трансформер, названная StyTr² Yingying Deng и ее сотрудниками. Эта модель содержит два разных кодировщика преобразователя для создания специфичных для предметной области последовательностей для содержимого и стиля соответственно. Затем эти кодировки вводятся в декодер многоуровневого преобразователя для стилизации последовательности контента в соответствии с последовательностью стилей. Преимущество преобразователей по сравнению с другими популярными методами переноса стилей, такими как генеративно-состязательные сети (GAN), заключается в том, что они очень быстро делают выводы. В конкретном случае нашего приложения преобразователь обрабатывает изображение за 2,52 секунды (на графическом процессоре NVIDIA Tesla K80)!
  • Вторая модель, которую мы используем, называется быстрая передача художественного стиля и легко доступна на Tensorflows model hub Голназа Гиази и соавторов. Модель расширяет первоначальную реализацию этой статьи В. Дюмулена и др. к произвольным стилям. Для этого они используют сеть прогнозирования стиля, которая использует архитектуру Inception-V3 для прогнозирования встроенного вектора изображения стиля. Эта модель в 2 раза быстрее в выводе, чем трансформер, но соответственно страдает правдоподобность результата, так как модель часто теряет трехмерность дивана.
  • Третью модель, которую мы реализовали, можно найти на PaddleHub и она называется stylepro Artistic. Модель основана на архитектуре StyleProNet, в которой используется новый метод преобразования стиля на уровне функций, названный Style Projection, для быстрого и эффективного преобразования стиля содержимого без параметров. Дополнительно реализация позволяет взвешивать стиль и оригинальное изображение. Эта модель предлагает более реалистичные результаты, чем предыдущая модель, но обработка изображения занимает примерно 7 секунд.

Конечно, для этой задачи существует множество других интересных моделей. Несколько примеров включают кодовую базу Neural Style Transfer Сомшубры Маджумдара и латентный редактор styleGAN ぱしふぃん, которые оба были отложены из-за более длительного времени выполнения.

Как я общаюсь с людьми?

Если мы хотим сделать помощника ИИ доступным для широкой публики, всегда приятно иметь опрятный интерфейс! Не имея опыта работы с HTML или CSS, поначалу это казалось сложной задачей. К счастью, мы нашли несколько простых в использовании пакетов Python с именами Gradio и Streamlit.

Я начал с Градио. Ниже вы можете увидеть общую структуру кода, который мне нужен для этого внешнего интерфейса.

import gradio as gr

def style_sofa( Input_image, Style_image, Choice_of_algorithm):

    # Preprocess input images
    resized_img, box = resize_sofa(fix_orient(Input_image))
    resized_style = resize_style(fix_orient(Style_image))

    # Generate mask for image
    mask = get_mask(resized_img)

    # Created a styled sofa
    styled_sofa = create_styledSofa(resized_img, resized_style, Choice_of_algorithm)
    
    # Postprocess the final image
    new_sofa = replace_sofa(resized_img, mask, styled_sofa)
    new_sofa = new_sofa.crop(box)
    
    return new_sofa
    
    
    
demo = gr.Interface(
    style_sofa,
    inputs=[
        gr.inputs.Image(type="pil"),
        gr.inputs.Image(type="pil"),
        gr.inputs.Radio(
            ["Style Transformer", "Style FAST", "Style Projection"],
            default="Style FAST",
        ),
    ],
    outputs="image",
    title="🛋 Style your sofa 🛋 ",
    description="Customize your sofa to your wildest dreams 💭!\
                \nProvide a picture of your sofa, a desired pattern\
                and (optionally) choose one of the algorithms.\
                \nOr just pick one of the examples below. ⬇",
    theme="huggingface",
    enable_queue=True,
)

if __name__ == "__main__":
    demo.launch(cache_examples=True)

Это приводит к внешнему интерфейсу, как показано в демо ниже:

0:00

/

1×

Демонстрация внешнего интерфейса, который я сделал с помощью Gradio

Однако позже я перешел на Streamlit. Оба пакета схожи по функциональности, но я обнаружил, что Streamlit обеспечивает большую гибкость в общем расположении приложения. Ниже вы можете увидеть новый и улучшенный интерфейс, который я сделал:

0:00

/

1×

Демонстрация интерфейса, который я сделал с помощью Streamlit

Размещение приложения на моем собственном доисторическом ноутбуке, конечно же, никогда не было целью. К счастью, я наткнулся на HuggingFace Spaces, который позволяет мне постоянно размещать мое приложение бесплатно! Интересно, как я это сделал? Проверьте это в моем другом блоге!

С моделью сегментации, переносом стилей и внешним интерфейсом у нас теперь есть готовый и работающий ИИ-помощник! И хотя его навыки сегментации великолепны, ему предстоит многому научиться, прежде чем он будет готов кому-то помочь, но я верю, что SofaStyler может спасти мир! …. или хотя бы свою гостиную.