Written by Aleksandra Kos• 27 sierpnia, 2021• 10:00 am• Deep Learning, Tutoriale, Video

SOTR, czyli segmentacja instancji z wykorzystaniem transformera

Mniej więcej rok temu opisywałam na blogu w jaki sposób uruchomić inferencję modelu DETR do wykrywania obiektów i segmentacji instancji. Czemu o tym wspominam? Bardzo niedawno ukazał się SOTR, czyli Segmenting Objects with Transformer, czyli kolejny model, który z powodzeniem wykorzystuje znanego z przetwarzania języka naturalnego (NLP) transformera w zagadnieniach analizy obrazu. Jeśli ciekawi Cię w jaki sposób transformer działa, tutaj możesz przeczytac więcej na jego temat.

Plan działania i organizacja

Na githubie deepdrive.pl znajdziesz repozytorium SOTR, w którym znajduje się skrypt do inferencji. Jest to fork oryginalnego repo z niewielkimi zmianami (o nich nieco później).
Posiłkować będziemy się tym samym filmem, który już przewinął się kilka razy, czy to na tym blogu, czy na kanale YouTube. Możesz go objerzeć tutaj.

Co chcemy uzyskać? Wyjściem będzie film z wizualizacją detekcji (masek obiektów).

Przygotowanie, czyli instalacje

W pierwszej kolejności stworzymy środowisko w anacondzie i w nim będziemy instalowac wszystkie potrzebne paczki. Poniższy kod tworzy i aktywuje środowisko.

conda create -n SOTR python=3.6
conda activate SOTR

W przypadku PyTorcha ważne jest jaką wersję CUDA mamy zainstalowaną (o instalacji trochę wspominał Karol przy okazji posta o budowaniu YOLOv4).

Załóżmy jednak, że CUDA jest u nas zainstalowana. Jak sprawdzić, którą wersję mamy. Można to zrobić na kilka sposóbów. Ja skorzystałam z poniższego polecenia.

cat /usr/local/cuda/version.txt

Po uzyskaniu informacji i wersji (w moim przypadku jest to 10.2), musimy wybrać się na stronę pytorcha i znaleźć instrukcję do instalacji ospowiedniej wersji. U mnie wystarczy poniższe.

pip install torch torchvision

Potrzebny nam także będzie detectron2 – instrukcję znajdziesz w README. Ponieważ mam CUDA 10.2 i PyTorcha 1.9, instalację wykonam w następujący sposób.

python -m pip install detectron2 -f   https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html

Doinstaluję jeszcze OpenCV i gdown, które pozwala na wygodne konsolowe pobieranie plików z Google Drive.

pip install opencv-python gdown

Wreszcie możemy przejść do repo SOTR.

git clone https://github.com/deepdrivepl/SOTR.git
cd SOTR

Pobieranie plików

Przydadzą nam się pretrenowane modele…

mkdir models
gdown 'https://drive.google.com/u/0/uc?export=download&confirm=taWO&id=1CzQTsvn9vxLnFkDJpIlitFXu1X_vw1dZ' -O models/SOTR_R101.pth
gdown 'https://drive.google.com/u/0/uc?id=19Dy6sXrwaNwGwNvuQyv5pZMWGM_at0ym&export=download' -O models/SOTR_R101_DCN.pth

…oraz wspomniany wcześniej film.

mkdir data
cd data
wget https://archive.org/download/0002201705192/0002-20170519-2.mp4

Na tym etapie, w głównym katalogu powinny znaleźć się dwa nowe podkatalogi.

models/
├── SOTR_R101_DCN.pth
└── SOTR_R101.pth
data/
└── 0002-20170519-2.mp4

Inferencja krok po kroku

We wspomnianym wcześniej repozytorium znajdziesz skrypt SOTR_inference.py. Pisząc go opierałam się na kodzie z katalogu demo. Pierwotnie miał to być notebook, ale z racji tego, że uruchamiałam inferencję dla dwóch różnych modeli tak było wygodniej (no i bardziej schludnie).

Importy

W pierwszej kolejności importy – zacznijmy od tych standardowych.

import os
import argparse
import time

from tqdm import tqdm

import cv2
import torch
import numpy as np

Następnie przyda nam się jeszcze kilka rzeczy z repo SOTR/detectrona2.

from adet.config import get_cfg
from adet.utils.visualizer import TextVisualizer
from demo.predictor import VisualizationDemo
from detectron2.utils.visualizer import Visualizer

Wizualizacja

Jeśli spojrzysz na klasę VisualizationDemo, zauważysz, że mamy praktycznie gotowy kod do wykonania inferencji i wizualizacji – metoda run_on_image() zwraca nam predykcję oraz zwizualizowane na oryginalnym obrazie detekcje. Brakowało mi dwóch rzeczy – możliwości odfiltrowania wyników o niskim score oraz informacji o czasie inferencji, który chciałam umieścić na wyjściowym filmie – dodałam je w klasie Visualization dziedziczącej po VisualizationDemo.

class Visualization(VisualizationDemo):
    def __init__(self, cfg, score_threshold=0.3):
        super(Visualization, self).__init__(cfg)
        self.score_threshold = score_threshold

    def run_on_image(self, image):
        vis_output = None
        t0=time.time()
        predictions = self.predictor(image)
        t1=time.time()
        # Convert image from OpenCV BGR format to Matplotlib RGB format.
        image = image[:, :, ::-1]
        if self.vis_text:
            visualizer = TextVisualizer(image, self.metadata, instance_mode=self.instance_mode)
        else:
            visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode)

        if "instances" in predictions:
            instances = predictions["instances"].to(self.cpu_device)
            instances = instances[instances.scores >= self.score_threshold]
            vis_output = visualizer.draw_instance_predictions(predictions=instances)

        return instances, vis_output, t1-t0

Kluczowe są 3 różnice: pomiar i zwrócenie czasu inferencji, dodanie self.score_threshold oraz odfiltrowanie predykcji instances = instances[instances.scores >= self.score_threshold].

W pierwotnej wersji zrobiłam tak jak powyżej, czyli wykorzystałam metodę draw_instance_predictions(). W przypadku pracy z filmami nie zdaje ona egzaminu, ponieważ każda instancja klasy jest oznaczana losową barwą. Skutek? Ten sam obiekt na kolejnych klatkach ma zupełnie inne kolory. Przez to w filmie powinno znaleźć się ostrzeżenie przed atakiem epilepsji.

Pierwotnie było tak:

vis_output = visualizer.draw_instance_predictions(predictions=instances)

Ostatecznie powyższy kod zmieniłam na:

labels = [self.metadata.thing_classes[x] for x in instances.pred_classes]
colors = [self.metadata.thing_colors[x] for x in instances.pred_classes]
colors = [[x/255 for x in lst] for lst in colors]
vis_output = visualizer.overlay_instances(boxes=instances.pred_boxes, labels=labels, masks=instances.pred_masks, assigned_colors=colors, alpha=0.3)

Musimy kilka rzeczy zrobić na piechotę, w efekcie mamy jednak zapewnione, że obiekty należące do tej samej klasy będą miały zawsze jednakowy kolor.

Wczytanie modelu

Kolejna funkcja ogarnia nam wczytanie pliku konfiguracyjnego, który jest podstawą modeli zaimplementowanych z wykorzystaniem detectrona2.

def setup_cfg(model_cfg, model_path):
    cfg = get_cfg()
    cfg.merge_from_file(model_cfg)
    cfg.MODEL.WEIGHTS = model_path
    cfg.freeze()
    return cfg

Jak zapewne udało Ci się zauważyć, zwracany plik konfiguracyjny zawiera już model z wczytanymi wagami. W zasadzie mamy teraz już wszystko, czego potrzebujemy do przetworzenia naszego filmu.

Segmentacja instancji na filmie

def main(video_path, out_dir, demo, model_name):
    cap = cv2.VideoCapture(video_path)

    idx = 0
    pbar = tqdm(total=int(cap.get(cv2.CAP_PROP_FRAME_COUNT)))
    while(cap.isOpened()):
        ret, frame = cap.read()
        if frame is None:
            break
        
        predictions, visualized_output, inf_time = demo.run_on_image(frame)
        visualized_output = visualized_output.get_image()
        
        txt="%s Inference: %dx%d  GPU: %s Inference time %.3fs" % (model_name,750,1333,torch.cuda.get_device_name(0), inf_time)
        cv2.putText(visualized_output,txt, (100,100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,0),19)
        cv2.putText(visualized_output,txt, (100,100), cv2.FONT_HERSHEY_SIMPLEX, 2, (255,255,255),9)

        cv2.imwrite(os.path.join(out_dir, 'img%08d.jpg' % idx), visualized_output[:,:,::-1])
        idx+=1; pbar.update(1)
        del frame

    cap.release()

Powyższy kod odczytuje film klatka po klatce – wykonuje inferencję i wizualizację oraz dodaje interesujące nas statystyki, czyli nazwę modelu, rozdzielczość, model GPU oraz czas inferencji.

Jedna rzecz w powyższym kodzie jest zrobiona mało elegancko – jest to wpisana na sztywno rozdzielczość. Po prostu sprawdziłam do jakiej rozdzielczości SOTR zmniejsza zdjęcia w preprocessingu i ją tu wpisałam. Teoretycznie powinnam te wartości zwrócić z fukcji, która wykonuje przetwarzanie wstępne, ale wymagałoby to zmian w kodzie. Na potrzeby szybkiej wizualizacji jest wystarczająco dobrze :).

Opakowanie

No i na koniec opakowanie wszystkiego, czyli __main__.

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_cfg', type=str, required=True)
    parser.add_argument('--model_path', type=str, required=True)
    parser.add_argument('--video_path', type=str, required=True)
    parser.add_argument('--out_dir', type=str, required=True)
    parser.add_argument('--score_threshold', type=int, default=0.2)
    args = parser.parse_args()
    
    if not os.path.exists(args.out_dir):
        os.makedirs(args.out_dir)
        
    cfg = setup_cfg(args.model_cfg, args.model_path)
    demo = Visualization(cfg, score_threshold=args.score_threshold)
    
    model_name = os.path.splitext(os.path.basename(args.model_path))[0]
    main(args.video_path, args.out_dir, demo, model_name)

Skrypt można uruchomić poniższym poleceniem.
python SOTR_inference.py --model_cfg configs/SOTR/R101.yaml --model_path models/SOTR_R101.pth --video_path data/0002-20170519-2.mp4 --out_dir data/results/SOTR_R101
Jeśli chcesz zmienić score_threshold (na przykład na 0.5) możesz dodać także --score_threshold 0.5.

Ciekawostka

Przy pierwszym uruchomieniu skryptu, etykiety nie wyświetlały się prawidłowo, dokładniej wszystkie znalazły się w lewym górnym rogu. Położenie etykiet wyznaczanie jest na podstawie bounding-boxów, których obliczanie zostało w repozytorium SOTR zakomentowane – odkomentowanie załatwiło sprawę i wizualizacja zaczęła wyglądać dobrze.

Powyższy komentarz nie wymaga z Twojej strony żadnego działania, ponieważ opisana zmiana jest w repozytorium.

Generowanie filmu

W wyniku wykonania skryptu SOTR_inference.py uzyskamy zwizualizowane kolejne klatki. W uzyskaniu filmu pomoże nam ffmpeg.
ffmpeg -i data/results/SOTR_R101/img%08d.jpg data/results/SOTR_R101.mp4

Wynik

SOTR R101

Podsumowanie

Na koniec chciałabym wspomnieć o raportowanej średniej precyzji, którą SOTR osiąga na zbiorze COCO test-dev. Jeśli spojrzysz na ranking, SOTR znajduje się dość daleko (na moment pisania posta zajmuje 17 miejsce). Jeśli przyjrzymy się jednak dokładniej wynikom, możemy zauważyć, że średnia precyzja dla dużych i średnich obiektów jest wysoka (obecnie 1 miejsce APL i APM). Niska jest natomiast wartość APS (zaledwie 11.5).

Warto mieć to na uwadze, zanim zdecydujemy się na wykorzystanie tego modelu w jakimś swoim projekcie – jeśli wiemy, że będą tam małe obiekty, powinniśmy prawdopodobnie zdecydować się na inny model.