Hai semuanya! Saya seorang insinyur penelitian di tim visi komputer Grup Mail.ru. Dalam artikel ini, saya akan menceritakan tentang bagaimana kami telah menciptakan proyek pemulihan foto berbasis AI untuk foto-foto militer lama. Apa itu "restorasi foto"? Ini terdiri dari tiga langkah:

• kami menemukan semua cacat gambar: fraktur, pertengkaran, lubang;
  
• kami mengecat cacat yang ditemukan, berdasarkan nilai piksel di sekitarnya;
  
• kami mewarnai gambar.
  

Lebih lanjut, saya akan menjelaskan setiap langkah pemulihan foto dan memberi tahu Anda bagaimana kami mendapatkan data kami, jaring apa yang kami latih, apa yang kami capai, dan kesalahan apa yang kami buat.

Mencari cacat

Kami ingin menemukan semua piksel yang terkait dengan cacat pada foto yang diunggah. Pertama, kita perlu mencari tahu gambar seperti apa yang akan diunggah orang. Kami berbicara dengan para pendiri proyek "Immortal Regiment", sebuah organisasi non-komersial yang menyimpan foto-foto warisan WW2, yang membagikan data mereka kepada kami. Setelah menganalisisnya, kami perhatikan bahwa orang-orang mengunggah sebagian besar potret individu atau kelompok dengan cacat tingkat sedang hingga besar.

Kemudian kami harus mengumpulkan satu set pelatihan. Set pelatihan untuk tugas segmentasi adalah gambar dan topeng tempat semua cacat ditandai. Cara termudah untuk melakukannya adalah membiarkan penilai membuat topeng segmentasi. Tentu saja, orang tahu betul cara menemukan cacat, tetapi itu akan memakan waktu terlalu lama.



Butuh satu jam atau seluruh hari kerja untuk menandai piksel cacat dalam satu foto. Karena itu, tidak mudah untuk mengumpulkan satu set pelatihan lebih dari 100 gambar dalam beberapa minggu. Itu sebabnya kami mencoba menambah data kami dan membuat cacat kami sendiri: kami akan mengambil foto yang bagus, menambahkan cacat menggunakan jalan acak pada gambar, dan berakhir dengan topeng yang menunjukkan bagian gambar dengan cacat. Tanpa augmentasi, kami memiliki 68 foto berlabel secara manual di set latihan dan 11 foto di set validasi.

Pendekatan segmentasi yang paling populer: ambil Unet dengan encoder pra-terlatih dan minimalkan jumlah BCE ( binary cross-entropy ) dan DICE ( Sørensen - Dice koefisien ).

Masalah apa yang muncul ketika kita menggunakan pendekatan segmentasi ini untuk tugas kita?

• Bahkan jika terlihat ada banyak cacat dalam foto, bahwa itu sangat tua dan lusuh, area dengan cacat masih jauh lebih kecil daripada yang tidak rusak. Untuk mengatasi masalah ini, kita dapat meningkatkan bobot kelas positif dalam BCE; bobot yang optimal adalah rasio piksel bersih dengan yang rusak.
  
• Masalah kedua adalah bahwa jika kita menggunakan Unet out-of-box dengan pra-terlatih encoder (Albunet-18, misalnya), kita kehilangan banyak data posisi. Lapisan pertama Albunet-18 terdiri dari konvolusi dengan kernel 5 dan langkah yang sama dengan dua. Ini memungkinkan internet bekerja cepat. Kami menukar waktu operasi bersih untuk memiliki lokalisasi cacat yang lebih baik: kami menghapus kumpulan maksimum setelah lapisan pertama, mengurangi langkah ke 1 dan mengurangi konvolusi kernel menjadi 3.
  
• Jika kita bekerja dengan gambar kecil dengan mengompresnya, misalnya menjadi 256 x 256 atau 512 x 512 piksel, maka cacat kecil akan hilang karena interpolasi. Karena itu, kita perlu bekerja dengan gambar yang lebih besar. Kami saat ini mengelompokkan cacat dalam ukuran foto 1024 x 1024 dalam produksi. Itu sebabnya kami harus melatih jaring pada tanaman gambar besar. Namun, ini menyebabkan masalah dengan ukuran batch kecil pada satu GPU.
  
• Selama pelatihan, kami dapat memuat sekitar 20 gambar pada satu GPU. Karena itu, kita berakhir dengan nilai mean dan standar deviasi yang tidak akurat di lapisan BatchNorm. Kita dapat mengatasi masalah ini menggunakan BatchNorm Di Tempat , yang, di satu sisi, menghemat ruang memori, dan di sisi lain, memiliki versi BatchNorm yang Disinkronkan, yang menyinkronkan statistik di semua GPU. Sekarang kita menghitung nilai rata-rata dan standar deviasi bukan untuk 20 gambar pada satu GPU, tetapi untuk 80 gambar dari 4 GPU. Ini meningkatkan konvergensi bersih.
  

Akhirnya, setelah menambah bobot BCE, mengubah arsitektur, dan menggunakan BatchNorm di tempat, kami membuat segmentasi lebih baik. Namun, tidak perlu terlalu banyak biaya untuk melakukan sesuatu yang lebih baik dengan menambahkan Augmentasi Waktu Tes. Kita dapat menjalankan net sekali pada gambar input, lalu mirror itu dan jalankan kembali net untuk menemukan semua cacat kecil.



Jaring bertemu dalam 18 jam pada empat GeForce 1080Ti. Inferensi membutuhkan waktu 290 ms. Ini cukup panjang, tapi itulah harga kinerja kami yang lebih baik daripada standar. Validasi DICE sama dengan 0,35, dan ROCAUC - 0,93.

Pengecatan gambar

Sama dengan tugas segmentasi kami menggunakan Unet. Untuk melakukan pewarnaan, kami akan mengunggah gambar asli dan topeng di mana kami menandai semua area bersih dengan yang, dan dengan nol - semua piksel yang ingin kami cat. Ini adalah cara kami mengumpulkan data: untuk foto apa pun dari dataset gambar sumber terbuka, misalnya, OpenImagesV4, kami menambahkan cacat yang mirip dengan yang kita lihat dalam kehidupan nyata. Kemudian kami melatih jaring untuk mengembalikan bagian yang hilang.

Bagaimana kita bisa memodifikasi Unet untuk tugas ini?

Kita dapat menggunakan konvolusi parsial alih-alih yang asli. Idenya adalah bahwa ketika kita menggabungkan area dengan beberapa kernel, kita tidak memperhitungkan nilai piksel cacat. Ini membuat pewarnaan lebih tepat. Kami menunjukkan kepada Anda sebuah contoh dari makalah NVIDIA terbaru. Mereka menggunakan Unet dengan konvolusi 2 dimensi default di gambar tengah dan konvolusi parsial - dalam gambar di sebelah kanan.



Kami melatih jaring selama lima hari. Pada hari terakhir, kami membekukan BatchNorms untuk membuat batas bagian yang dicat tidak terlihat.

Dibutuhkan net 50 ms untuk memproses satu gambar 512 x 512. Validasi PSNR sama dengan 26,4. Namun, Anda tidak dapat sepenuhnya mengandalkan metrik dalam tugas ini. Untuk memilih model terbaik, kami menjalankan beberapa model bagus pada gambar penilaian, menganonimkan hasilnya, dan memilih yang paling kami sukai. Itulah cara kami memilih model akhir kami.

Saya telah menyebutkan sebelumnya bahwa kami secara artifisial menambahkan beberapa cacat pada gambar bersih. Anda harus selalu melacak ukuran maksimum cacat tambahan selama pelatihan; dalam kasus ketika Anda memberi makan gambar dengan cacat sangat besar ke jaring yang tidak pernah ditangani pada tahap pelatihan, jaring akan menjadi liar dan menghasilkan hasil yang tidak dapat diterapkan. Karena itu, jika Anda perlu memperbaiki cacat yang besar, tambah set latihan Anda dengan mereka.

Berikut adalah contoh cara kerja algoritma kami:

Pewarnaan

Kami mengelompokkan cacat dan mencatnya; langkah ketiga - rekonstruksi warna. Seperti yang saya katakan sebelumnya, ada banyak potret individu dan kelompok di antara foto-foto Resimen Abadi. Kami ingin jaring kami bekerja dengan baik bersama mereka. Kami memutuskan untuk membuat pewarnaan sendiri karena tidak ada layanan yang ada yang dapat mewarnai potret dengan cepat dan efisien. Kami ingin foto berwarna kami menjadi lebih dapat dipercaya.



GitHub memiliki repositori populer untuk pewarnaan foto. Itu pekerjaan yang baik tetapi masih memiliki beberapa masalah. Misalnya, ia cenderung mengecat pakaian biru. Itu sebabnya kami menolaknya juga.

Jadi, kami memutuskan untuk membuat algoritma untuk pewarnaan gambar. Ide paling jelas: ambil gambar hitam-putih dan prediksi tiga saluran: merah, hijau, dan biru. Namun, kita dapat membuat pekerjaan kita lebih mudah: bekerja bukan dengan representasi warna RGB, tetapi dengan representasi warna YCbCr. Komponen Y adalah kecerahan (luma). Gambar hitam-putih yang diunggah adalah saluran Y, dan kami akan menggunakannya kembali. Sekarang kita perlu memprediksi Cb dan Cr: Cb adalah perbedaan warna biru dan kecerahan dan Cr - perbedaan warna merah dan kecerahan.



Mengapa kami memilih representasi YCbCr? Mata manusia lebih sensitif terhadap perubahan kecerahan daripada perubahan warna. Itu sebabnya kami menggunakan kembali komponen Y (kecerahan) yang mata manusia paling peka terhadap dan memprediksi Cb dan Cr yang mungkin kita buat salah karena kita tidak bisa melihat kepalsuan warna dengan sangat baik. Karakteristik khusus ini banyak digunakan pada awal televisi berwarna ketika kapasitas saluran tidak cukup untuk mengirimkan semua warna. Gambar ditransmisikan dalam YCbCr, tidak berubah ke komponen Y, dan Cb dan Cr berkurang setengahnya.

Cara membuat garis dasar

Kita dapat mengambil Unet dengan encoder pretrained dan meminimalkan Kehilangan L1 antara nilai CbCr yang ada dan yang diprediksi. Kami ingin mewarnai potret dan, oleh karena itu, selain foto OpenImages, kami membutuhkan lebih banyak foto spesifik tugas.

Di mana kita bisa mendapatkan foto berwarna orang yang mengenakan seragam militer? Ada orang di internet yang mewarnai foto-foto lama sebagai hobi atau untuk harga. Mereka melakukannya dengan sangat hati-hati, berusaha menjadi sangat tepat. Ketika mereka mewarnai seragam, papan bahu, dan medali, mereka merujuk pada bahan arsip, sehingga hasil pekerjaan mereka dapat dipercaya. Secara keseluruhan, kami menggunakan 200 gambar berwarna secara manual dengan orang-orang berseragam militer.

Sumber data lain yang bermanfaat adalah situs web Tentara Merah Pekerja dan Petani . Salah satu pendirinya berfoto di hampir setiap seragam Soviet Perang Dunia 2 yang tersedia.



Dalam beberapa gambar, ia meniru pose orang-orang dari foto arsip terkenal. Ini adalah hal yang baik bahwa gambar-gambarnya memiliki latar belakang putih: itu memungkinkan kami untuk menambah data dengan sangat baik dengan menambahkan berbagai objek alami di latar belakang. Kami juga menggunakan beberapa potret reguler, melengkapi mereka dengan lencana dan atribut perang lainnya.

Kami melatih AlbuNet-50 - ini adalah Unet yang menggunakan ResNet-50 yang dipra-pretr sebagai encoder. Jaring mulai memberikan hasil yang memadai: kulit berwarna merah muda, mata - abu-abu-biru, papan bahu - kekuningan. Namun, masalahnya adalah bahwa beberapa area pada foto tidak tersentuh. Hal ini disebabkan oleh fakta bahwa menurut kesalahan L1 menemukan yang optimal di mana lebih baik untuk tidak melakukan apa-apa daripada mencoba memprediksi beberapa warna.


Kami membandingkan hasil kami dengan foto Ground Truth - pewarnaan manual yang dilakukan oleh Klimbim

Bagaimana kita bisa menyelesaikan masalah ini? Kita membutuhkan pembeda: jaringan saraf yang akan menerima gambar dan memberi tahu kita apakah itu terlihat realistis atau tidak. Salah satu gambar di bawah ini diwarnai secara manual dan yang lainnya - oleh generator kami, AlbuNet-50. Bagaimana manusia membedakan foto yang diwarnai secara manual dan otomatis? Dengan melihat detail. Bisakah Anda tahu di mana foto yang diwarnai secara otomatis dengan solusi dasar kami?




Kami menggunakan diskriminator dari kertas Self-Attention GAN . Ini adalah jaring konvolusi kecil dengan apa yang disebut Self-Attention yang dibangun di lapisan atas. Ini memungkinkan kita untuk "lebih memperhatikan" detail gambar. Kami juga menggunakan normalisasi spektral. Anda dapat menemukan informasi lebih lanjut di makalah yang disebutkan di atas. Kami telah melatih internet dengan kombinasi kehilangan L1 dan kehilangan dari diskriminator. Sekarang jaring mewarnai detail gambar dengan lebih baik, dan latar belakang terlihat lebih konsisten. Satu contoh lagi: di sebelah kiri adalah pekerjaan yang dilatih hanya dengan kehilangan L1; di sebelah kanan - dengan kombinasi kerugian diskriminator L1.



Proses pelatihan memakan waktu dua hari pada empat GeForce 1080Ti. Dibutuhkan 30 ms bersih untuk memproses gambar 512 x 512. Validasi MSE - 34,4. Sama seperti dengan pewarnaan, metrik yang tidak ingin Anda andalkan pada metrik. Itu sebabnya kami memilih enam model dengan metrik validasi terbaik dan memilih secara buta untuk model terbaik.

Ketika kami telah membuat sistem produksi dan meluncurkan situs web, kami terus bereksperimen dan menyimpulkan bahwa lebih baik kami meminimalkan bukan kerugian L1 per-piksel, tetapi kerugian persepsi. Untuk menghitungnya, kami mengumpankan prediksi netto dan foto ground-truthl ke jaring VGG-16, mengambil peta fitur di lapisan bawah dan membandingkannya dengan MSE. Pendekatan ini melukis lebih banyak area dan memberikan hasil yang lebih berwarna.

Rekap

Unet adalah model yang sangat keren. Pada tugas segmentasi pertama, kami menghadapi masalah selama pelatihan, dan bekerja dengan gambar resolusi tinggi dan itulah sebabnya kami menggunakan BatchNorm In-Place. Pada tugas kedua kami (Inpainting), kami menggunakan Konvolusi Parsial alih-alih yang default, dan memungkinkan kami untuk mendapatkan hasil yang lebih baik. Saat mengerjakan pewarnaan, kami menambahkan jaring pembeda kecil yang menghukum generator untuk gambar yang tidak realistis. Kami juga menggunakan kerugian persepsi.

Kesimpulan kedua - penilai sangat penting. Dan tidak hanya selama tahap pembuatan mask segmentasi tetapi juga untuk validasi hasil akhir. Pada akhirnya, kami memberikan pengguna tiga foto: gambar asli dengan cacat inpainted, foto berwarna dengan cacat inpainted dan yang hanya berwarna jika algoritma untuk pencarian cacat dan inpainting salah.

Kami mengambil beberapa gambar dari proyek Album Perang dan memprosesnya melalui neuronet ini. Inilah hasil yang kami dapatkan:



Selain itu, di sini Anda dapat melihat lebih dekat pada gambar asli dan semua tahapan pemrosesan.