Jak nauczyć robota widzieć obiekty bez długiego treningu

Wyobraź sobie: kupiłeś nowy manipulator lub postanowiłeś zbudować inteligentną kamerę do sortowania części. Masz mnóstwo różnych obiektów — od puszek po napoje po złożone części zamienne. Standardowa ścieżka w Computer Vision dzisiaj to: zbierz dataset, oznakuj tysiące klatek, wytrenuj model i miej nadzieję, że nie będzie „dryfował" pod różnym oświetleniem. A jeśli jutro pojawi się nowy obiekt? Będziesz musiał zacząć wszystko od nowa.

Inżynierowie z NVIDIA Research najwyraźniej zmęczyli się tym niekończącym się cyklem i wypuścili FoundationPose. To ujednolicony model do estymacji pozy 6D (pozycja i orientacja w przestrzeni) oraz śledzenia obiektów, który działa „out-of-the-box" nawet z obiektami, których nigdy nie widział podczas treningu.

Na czym polega główna feta

Zazwyczaj zadania estymacji pozy 6D dzielą się na dwa obozy. W pierwszym masz precyzyjny model CAD obiektu (podejście model-based), a w drugim masz tylko kilka zdjęć z różnych kątów (podejście model-free). FoundationPose łączy te podejścia.

Projekt jest interesujący, ponieważ nie wymaga fine-tuningu pod kątem konkretnego zadania. Wystarczy podać mu model 3D lub zestaw obrazów referencyjnych, a on natychmiast zaczyna wyprowadzać współrzędne obiektu w przestrzeni. Obecnie to rozwiązanie zajmuje pierwsze miejsce w światowym rankingu BOP (Benchmark for 6D Object Pose Estimation) wśród metod dla nowych obiektów.

Jak to działa pod maską

Deweloperzy zastosowali kilka sprytnych sztuczek, aby osiągnąć taką elastyczność.

Po pierwsze, użyli neural implicit representation. Pozwala to systemowi syntetyzować nowe widoki obiektu w locie, dzięki czemu proces estymacji pozy jest identyczny zarówno w przypadku gotowego modelu 3D, jak i kilku zdjęć.

Po drugie, skala treningu jest imponująca. Model był trenowany na ogromnej ilości danych syntetycznych. Ciekawy szczegół: duże modele językowe (LLM) były zaangażowane w generowanie tych danych. Pomogło to stworzyć zróżnicowane scenariusze i tekstury, które trudno uzyskać ręcznie.

Po trzecie, architektura opiera się na transformerach i contrastive learning. Upraszczając — model uczy się porównywać bieżącą klatkę z referencyjną i dokładnie rozumieć, jak obiekt jest obrócony, bez przywiązywania się do specyficznych cech „tej konkretnej wiertarki".

Co projekt potrafi w praktyce

Jeśli spojrzysz na demy w repozytorium, zobaczysz trzy główne scenariusze użycia:

1. Robotyka. Robot chwyta butelkę musztardy, która stale się porusza. System zdąrza przeliczyć pozę w czasie rzeczywistym, co jest krytyczne dla manipulatorów.
2. Augmented reality (AR). Obiekty wirtualne są „przyklejane" do prawdziwych obiektów bez jitteringu i przesunięć.
3. Praca ze złożonymi datasetami. Model dobrze radzi sobie z klasycznymi benchmarkami jak YCB-Video, gdzie jest dużo okluzji (gdy jeden obiekt blokuje drugi).

Przy okazji, dla tych którzy dbają o szybkość w produkcji, NVIDIA przygotowała wersję ROS z obsługą TensorRT. Daje to znaczny wzrost FPS dzięki optymalizacji GPU.

Jak uruchomić i potestować

Deployowanie takich projektów ML często zamienia się w dependency hell, ale tutaj autorzy oferują dwa rozsądne podejścia.

Najprostsze to Docker. Zespół przygotował obraz ze wszystkimi sztuczkami CUDA już skonfigurowanymi.

cd docker/
docker pull wenbowen123/foundationpose && docker tag wenbowen123/foundationpose foundationpose
bash docker/run_container.sh

Jeśli wolisz Conda, proces jest nieco bardziej skomplikowany, bo będziesz musiał ręcznie zbudować rozszerzenia C++ i CUDA. Będziesz musiał zainstalować PyTorch, PyTorch3D i NVDiffRast.

Po skonfigurowaniu uruchomienie dem jest proste:

python run_demo.py

Domyślnie skrypt pokaże, jak model śledzi butelkę musztardy z danych demo. Najpierw jest inicjalizacja (estymacja pozy na pierwszej klatce), a potem automatyczne przełączenie w tryb śledzenia.

Niuanse i ograniczenia

Nie spodziewaj się, że wszystko będzie działać idealnie na każdej karcie graficznej. Na przykład właściciele RTX 4090 będą musieli użyć niestandardowego obrazu Docker ze względu na specyfikę nowych kerneli CUDA.

Inny ważny punkt: ze względu na ograniczenia licencyjne Stable Diffusion (który był użyty do generowania niektórych tekstur do treningu), autorzy nie mogli udostępnić wag modelu trenowanych z augmentacją diffusion. Wersja publiczna jest nieco prostsza, więc dokładność może być marginalnie niższa niż podano w oficjalnej publikacji naukowej.

Komu się to przyda

FoundationPose to nie tylko kolejna sieć neuronowa spośród setek publikacji CVPR. To gotowy fundament dla tych, którzy pracują nad:

• Robotyką magazynową (gdzie trzeba szybko dodawać nowe produkty do systemu).
• Kontrolą jakości w produkcji.
• Tworzeniem interaktywnych aplikacji AR.

Jeśli musisz określić pozycję obiektów w przestrzeni i nie chcesz spędzać tygodni na zbieraniu datasetów dla każdego nowego „śrubeczka", to to repozytorium zdecydowanie warto dodać do zakładek. Projekt jest aktywny, intensywnie dyskutowany w Issues, a prowadzenie w BOP sugeruje, że dziś to jeden z najsilniejszych metod SOTA w swojej niszy.