rgimad / ISR_GUI

GUI application for inferencing IR super-resolution models and analyzing their quality
1 stars 0 forks source link
pyqt pyqt5 python pytorch roi super-resolution

Infrared Super Resolution

Infrared Super Resolution - приложение с графическим интерфейсом на PyQt, позволяющая в интерактивном режиме увеличивать разрешение ИК изображений и анализировать качество работы обученных ИИ моделей (FSRCNN, VDSR, EDSR, ESCPCN, SRGAN) на заданных фрагментах с помощью метрик SSIM и PSNR.

Screen 1 Screen 2

Для написания программы использовался стек технологий Python и библиотеки PyQt, OpenCV, numpy, scikit-image, фреймворк глубокого обучения PyTorch. Графический пользовательский интерфейс программы был спроектирован в визуальной среде QtDesigner и представляет собой файл формата ui, который загружается при запуске программы. Из компонентов интерфейса, предоставляемых библиотекой Qt используются QTabWidget, QPushButton, QGroupBox, QComboBox QLabel, QAction. В программе также используются свои графические компоненты ImageLabel и RoiImageLabel, унаследованные от компонента QLabel, предназначенного для размещения текста и изображений. Компонент ImageLabel добавляет к оригинальному QLabel дополнительную функциональность в виде кнопок «Копировать» и «Сохранить» в контекстном меню правой кнопки мыши, а компонент RoiImageLabel реализует возможность выбирать прямоугольную область интереса (ROI, Region of interest) определенного размера на загруженном в него изображении.

Для использования нейросетевых модулей в программе используется фреймворк PyTorch. Веса обученных моделей хранятся в файлах с расширением .pth.tar и называются следующим образом: например, файл с весами обученной модели EDSR с двухкратным увеличением называется edsr_ir_x2.pth.tar. Когда программе требуется та или иная модель, создается ее экземпляр и в нее загружаются веса из соответствующего файла. Для создания собственно экземпляров моделей, в которые будут загружаться веса, в программе содержатся модули c определениями классов используемых нейросетевых моделей, например, espcn_ir_model.py. Для инференса моделей используется видеоускоритель, поддерживающий технологию CUDA, если таковой доступен, либо в противном случае задействуется центральный процессор ПК.

В левом верхнем углу окна программы пользователь может выбрать используемую нейросетевую модель из списка доступных, а также коэффициент увеличения разрешения: 2 или 4 (по каждой стороне соответственно). В программе предусмотрено два режима работы: режим исследования и режим инференса (основной). В режиме исследования пользователь выбирает интересующий фрагмент (ROI, Region of Interest), после чего он уменьшается, в 2 или в 4 раза, в зависимости от выбранного коэффициента масштабирования, и затем восстанавливается до исходного разрешения с помощью бикубической интерполяции и выбранной модели соответственно. В исследовательском режиме можно проводить сравнение результата нейросетевых моделей, бикубической интерполяции и исходного фрагмента в понятной визуальной форме и в виде метрик качества PSNR и SSIM.

Основной режим работы программы – увеличение разрешения ИК изображения. Для этого пользователю необходимо выбрать модель из выпадающего списка и коэффициент увеличения разрешения, затем выбрать файл с изображением, нажав на кнопку «Выбрать изображение». После этого в панели «Super Resolution» загрузится результат работы модели. При дальнейшем переключении модели и коэффициента увеличения разрешения, результат на данной панели будет пересчитываться автоматически.