Hai semuanya, saya bekerja sebagai programmer riset di tim visi komputer Mail.ru Group. Untuk Hari Kemenangan tahun ini, kami memutuskan untuk membuat proyek untuk pemulihan foto-foto militer . Apa itu restorasi foto? Ini terdiri dari tiga tahap:

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

Dalam artikel ini, saya akan membahas setiap tahap pemulihan secara terperinci dan memberi tahu Anda bagaimana dan di mana kami mengambil data, jaringan apa yang kami pelajari, apa yang kami lakukan, garu apa yang kami injak.

Pencarian cacat

Kami ingin menemukan semua piksel yang terkait dengan cacat pada foto yang diunggah. Pertama, kita perlu memahami jenis foto tahun perang yang akan diunggah orang. Kami berpaling ke penyelenggara proyek Resimen Immortal, yang berbagi data dengan kami. Setelah menganalisisnya, kami perhatikan bahwa orang-orang sering mengunggah potret, tunggal atau grup, yang memiliki sejumlah cacat sedang atau besar.

Maka perlu untuk mengumpulkan sampel pelatihan. Sampel pelatihan untuk tugas segmentasi adalah gambar dan topeng tempat semua cacat ditandai. Cara termudah adalah memberi foto ke spidol ke spidol. Tentu saja, orang pandai menemukan cacat, tetapi masalahnya adalah bahwa markup adalah proses yang sangat panjang.



Diperlukan waktu dari satu jam hingga satu hari kerja untuk menandai piksel terkait cacat dalam satu foto, sehingga sulit untuk mengumpulkan sampel lebih dari 100 foto dalam beberapa minggu. Karena itu, kami mencoba menambah data kami dan menulis cacat pada diri kami sendiri: kami mengambil foto yang bersih, menerapkan cacat buatan padanya dan mendapatkan topeng yang menunjukkan bagian-bagian tertentu dari gambar yang rusak. Bagian utama dari sampel pelatihan kami adalah 79 foto yang ditandai secara manual, yang 11 ditransfer ke sampel uji.

Pendekatan yang paling populer untuk masalah segmentasi: ambil Unet dengan pembuat enkode yang sudah terlatih dan minimalkan jumlahnya $$$SM$ ( binary cross-entropy ) dan $$$DICE$ ( Sørensen - Koefisien dadu ).

Masalah apa yang muncul dengan pendekatan ini dalam masalah segmentasi cacat?

• Sekalipun tampak bagi kita bahwa ada banyak cacat dalam foto, bahwa itu sangat kotor dan sangat compang-camping oleh waktu, area yang ditempati oleh cacat masih jauh lebih kecil daripada bagian gambar yang tidak rusak. Untuk mengatasi masalah ini, Anda dapat menambah bobot kelas positif di $$$SM$ , dan bobot optimal akan menjadi rasio dari jumlah semua piksel murni dengan jumlah piksel yang dimiliki cacat.
  
• Masalah kedua adalah bahwa jika kita menggunakan Unet out of the box dengan pra-terlatih encoder, misalnya Albunet-18, maka kita kehilangan banyak informasi posisi. Lapisan pertama Albunet-18 terdiri dari konvolusi dengan inti 5 dan berjalan sama dengan dua. Ini memungkinkan jaringan bekerja dengan cepat. Kami mengorbankan waktu operasi jaringan untuk lokalisasi cacat yang lebih baik: kami menghapus kumpulan maksimum setelah lapisan pertama, mengurangi langkah ke 1 dan mengurangi inti konvolusi menjadi 3.
  
• Jika kita bekerja dengan gambar kecil, misalnya, mengompres gambar ke 256 x 256 atau 512 x 512, maka cacat kecil akan hilang begitu saja karena interpolasi. Karena itu, Anda perlu bekerja dengan gambar besar. Sekarang dalam produksi kami melakukan segmentasi cacat pada foto 1024 x 1024. Oleh karena itu, perlu untuk melatih jaringan saraf pada potongan besar gambar besar. Dan karena ini, ada masalah dengan ukuran kecil batch pada satu kartu video.
  
• Selama pelatihan, kami memiliki sekitar 20 gambar yang ditempatkan pada satu kartu. Karena itu, estimasi rata-rata dan varians di lapisan BatchNorm tidak akurat. BatchNorm di tempat membantu kami mengatasi masalah ini, yang, pertama, menghemat memori, dan kedua, ia memiliki versi BatchNorm yang Disinkronkan, yang menyinkronkan statistik antara semua kartu. Sekarang kami mempertimbangkan rata-rata dan varians bukan dengan 20 gambar pada satu kartu, tetapi dengan 80 gambar dari 4 kartu. Ini meningkatkan konvergensi jaringan.
  

Pada akhirnya, berat badan bertambah $$$SM$ Dengan mengubah arsitektur dan menggunakan BatchNorm Di Tempat, kami mulai mencari cacat pada foto. Tetapi dengan murah Anda bisa melakukan sedikit lebih baik dengan menambahkan Test Time Augmentation. Kita dapat menjalankan jaringan sekali di gambar input, lalu mirror itu dan jalankan jaringan lagi, ini dapat membantu kita menemukan cacat kecil.



Akibatnya, jaringan kami melakukan konvergensi pada empat GeForce 1080Ti dalam 18 jam. Inferensi membutuhkan waktu 290 ms. Ternyata cukup lama, tetapi ini adalah harga untuk fakta bahwa kami sedang mencari cacat kecil. Validasi $$$DICE$ sama dengan 0,35, dan $$$ROCAUC$ - 0,93.

Restorasi fragmen

Unet membantu kami memecahkan masalah ini lagi. Untuk input, kami memberinya gambar asli dan topeng di mana kami menandai ruang bersih dengan unit, dan piksel yang ingin kita lukis dengan nol. Kami mengumpulkan data sebagai berikut: kami mengambil set data besar dengan gambar dari Internet, misalnya, OpenImagesV4, dan cacat buatan yang dibuat serupa dengan yang ditemukan dalam kehidupan nyata. Dan setelah itu mereka melatih jaringan untuk memperbaiki bagian yang hilang.

Bagaimana kita bisa memodifikasi Unet untuk tugas ini?

Anda dapat menggunakan Konvolusi Parsial alih-alih konvolusi biasa. Idenya adalah bahwa ketika kita menciutkan suatu wilayah gambar dengan kernel, kita tidak memperhitungkan nilai piksel yang terkait dengan cacat. Ini membantu membuat lukisan lebih akurat. Contoh dari artikel NVIDIA . Dalam gambar pusat, mereka menggunakan Unet dengan konvolusi biasa, dan di sebelah kanan - dengan Konvolusi Parsial:



Kami melatih jaringan selama 5 hari. Pada hari terakhir, kami membekukan BatchNorm, ini membantu membuat batas bagian gambar yang dicat kurang terlihat.

Jaringan memproses gambar 512 x 512 dalam 50 ms. Validasi PSNR adalah 26,4. Namun, metrik tidak dapat dipercaya tanpa syarat dalam tugas ini. Oleh karena itu, kami pertama-tama menjalankan beberapa model bagus pada data kami, menganonimkan hasilnya, dan kemudian memilih model yang kami sukai. Jadi kami memilih model final.

Saya menyebutkan bahwa kami menambahkan cacat buatan untuk membersihkan gambar. Saat berlatih, Anda perlu memonitor dengan hati-hati ukuran maksimum cacat yang ditumpangkan, karena dengan cacat yang sangat besar yang belum pernah dilihat jaringan dalam proses pembelajaran, ia akan secara liar berfantasi dan memberikan hasil yang benar-benar tidak pantas. Jadi, jika Anda perlu mengecat cacat besar, juga menerapkan cacat besar selama pelatihan.

Berikut adalah contoh algoritme:

Mewarnai

Kami mengelompokkan cacat dan mengecatnya, langkah ketiga adalah rekonstruksi warna. Biarkan saya mengingatkan Anda bahwa di antara foto-foto "Resimen Abadi" ada banyak potret tunggal atau kelompok. Dan kami ingin jaringan kami bekerja dengan baik dengan mereka. Kami memutuskan untuk membuat pewarnaan sendiri, karena tidak ada layanan yang kami kenal yang melukis potret dengan cepat dan baik.



GitHub memiliki repositori populer untuk mewarnai foto. Rata-rata, dia melakukan pekerjaan ini dengan baik, tetapi dia memiliki beberapa masalah. Misalnya, ia suka melukis pakaian dengan warna biru. Karena itu, kami juga menolaknya.

Jadi, kami memutuskan untuk membuat jaringan saraf untuk pewarnaan. Ide yang paling jelas: ambil gambar hitam putih dan prediksi tiga saluran, merah, hijau dan biru. Tetapi, secara umum, kita dapat menyederhanakan pekerjaan kita. Kita dapat bekerja bukan dengan representasi warna RGB, tetapi dengan representasi YCbCr. Komponen Y adalah kecerahan (luma). Gambar hitam putih yang diunduh adalah saluran Y, kami akan menggunakannya kembali. Tetap untuk memprediksi Cb dan Cr: Cb adalah perbedaan dalam warna biru dan kecerahan, dan Cr adalah perbedaan dalam warna merah dan kecerahan.



Mengapa kami memilih tampilan YCbCr? Mata manusia lebih rentan terhadap perubahan kecerahan daripada perubahan warna. Oleh karena itu, kami menggunakan kembali komponen-Y (kecerahan), sesuatu yang pada awalnya mudah diterima mata, dan memprediksi Cb dan Cr, di mana kita dapat membuat sedikit lebih banyak kesalahan, karena orang-orang melihat lebih sedikit "salah" dalam warna. Fitur ini mulai aktif digunakan pada awal televisi berwarna, ketika bandwidth saluran tidak cukup untuk mengirimkan semua warna secara penuh. Gambar dipindahkan ke YCbCr, ditransfer ke komponen Y tidak berubah, dan Cb dan Cr dikompresi dua kali.

Cara merakit garis dasar

Anda dapat kembali menggunakan Unet dengan encoder pra-terlatih dan meminimalkan Kehilangan L1 antara CbCr nyata dan yang diprediksi. Kami ingin mewarnai potret, jadi selain foto dari OpenImages, kami perlu menambahkan foto khusus untuk tugas kami.

Di mana saya bisa mendapatkan foto berwarna orang-orang berseragam militer? Ada orang di internet yang melukis foto-foto lama sebagai hobi atau untuk memesan. Mereka melakukan ini dengan sangat hati-hati, berusaha untuk sepenuhnya mematuhi semua nuansa. Mewarnai seragam, tanda pangkat, medali, mereka beralih ke bahan arsip, sehingga hasil karya mereka dapat dipercaya. Secara total, kami menggunakan 200 foto yang dilukis dengan tangan. Sumber data kedua yang bermanfaat adalah situs Tentara Merah Buruh dan Petani . Salah satu penciptanya difoto di hampir semua varian seragam militer selama Perang Patriotik Besar.



Dalam beberapa foto, ia mengulangi pose orang-orang dari foto arsip terkenal. Sangat bagus dia memotret dengan latar belakang putih, ini memungkinkan kami untuk menambah data dengan sangat baik, menambahkan berbagai objek alami ke latar belakang. Kami juga menggunakan potret orang modern biasa, melengkapi mereka dengan lencana dan atribut pakaian perang lainnya.

Kami melatih AlbuNet-50 - ini adalah Unet, di mana AlbuNet-50 digunakan sebagai encoder. Jaringan mulai memberikan hasil yang memadai: kulit berwarna merah muda, mata biru abu-abu, tali bahu berwarna kekuningan. Tapi masalahnya adalah dia melukis gambar-gambar itu dengan bintik-bintik. Hal ini disebabkan oleh fakta bahwa dari sudut pandang kesalahan L1 terkadang lebih menguntungkan untuk tidak melakukan apa pun daripada mencoba memprediksi warna.


Kami membandingkan hasil kami dengan foto Ground Truth - pewarnaan manual dari artis dengan julukan Klimbim

Bagaimana cara mengatasi masalah ini? Kita membutuhkan pembeda: jaringan saraf, di mana kita akan memasok gambar ke input, dan itu akan mengatakan seberapa realistis gambar ini terlihat. Di bawah, satu foto dilukis dengan tangan, dan yang kedua oleh jaringan saraf. Menurutmu yang mana?




Sebagai diskriminator, kami menggunakan diskriminator dari artikel Self-Attention GAN . Ini adalah jaringan konvolusional kecil, di lapisan terakhir yang disebut Self-Attention. Ini memungkinkan Anda untuk lebih "memperhatikan" detail gambar. Kami juga menggunakan normalisasi spektral. Penjelasan dan motivasi yang tepat dapat ditemukan di artikel. Kami melatih jaringan dengan kombinasi kehilangan-L1 dan kesalahan yang dikembalikan oleh diskriminator. Sekarang jaringan mengecat detail gambar dengan lebih baik, dan latar belakang lebih konsisten. Contoh lain: di sebelah kiri adalah hasil dari jaringan hanya dilatih dengan kehilangan-L1, di sebelah kanan - dengan kehilangan-L1 dan kesalahan diskriminator.



Pada empat Geforce 1080Ti, pelatihan butuh dua hari. Jaringan bekerja dalam 30 ms dalam gambar 512 x 512. MSE validasi adalah 34,4. Seperti pada masalah pewarnaan, metrik tidak dapat sepenuhnya dipercaya. Karenanya, kami memilih 6 model yang memiliki metrik terbaik untuk validasi, dan memilih secara buta untuk model terbaik.

Setelah meluncurkan model dalam produksi, kami melanjutkan percobaan dan sampai pada kesimpulan bahwa lebih baik untuk meminimalkan bukan per-pixel L1-loss, tetapi kerugian perseptual. Untuk menghitungnya, Anda perlu menjalankan prediksi jaringan dan foto sumber melalui jaringan VGG-16, mengambil peta atribut pada lapisan bawah dan membandingkannya sesuai dengan MSE. Pendekatan ini melukis lebih banyak area dan membantu mendapatkan gambar yang lebih berwarna.

Kesimpulan dan Kesimpulan

Unet adalah model yang keren. Dalam masalah segmentasi pertama, kami mengalami masalah dalam pelatihan dan bekerja dengan gambar resolusi tinggi, jadi kami menggunakan BatchNorm Di Tempat. Dalam tugas kedua (Inpainting), alih-alih konvolusi biasa, kami menggunakan Konvolusi Parsial, ini membantu mencapai hasil yang lebih baik. Dalam masalah pewarnaan untuk Unet, kami menambahkan jaringan diskriminator kecil yang mendenda generator untuk gambar tampak tidak realistis dan menggunakan kerugian persepsi.

Kesimpulan kedua adalah bahwa pengakses itu penting. Dan tidak hanya pada tahap menandai gambar sebelum pelatihan, tetapi juga untuk memvalidasi hasil akhir, karena dalam masalah cacat lukisan atau pewarnaan, Anda masih perlu memvalidasi hasil dengan bantuan seseorang. Kami memberikan pengguna tiga foto: yang asli dengan cacat yang dihapus, diwarnai dengan cacat yang dihapus, dan hanya foto berwarna jika algoritma untuk mencari dan melukis cacat salah.

Kami mengambil beberapa foto proyek Album Militer dan memprosesnya dengan jaringan saraf kami. Inilah hasil yang didapat:



Dan di sini Anda dapat melihatnya dalam resolusi asli dan pada setiap tahap pemrosesan.