bulk-image-tagger/application/main_window.py

# -*- coding: utf-8 -*-
import os
from glob import glob
from io import BytesIO

import cv2
import numpy as np
from PIL import Image
from PIL.ImageQt import ImageQt
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtGui import QPixmap
from PyQt5.QtWidgets import QFileDialog, QLabel


# Form implementation generated from reading ui file 'Bulk Image Tagger qt5.ui'
#
# Created by: PyQt5 UI code generator 5.14.1

def get_folder(folder):
    glob_pattern = sorted(glob(os.path.join(folder, '*')), key=os.path.getctime)
    out = []
    for i in range(len(glob_pattern)):
        if os.path.splitext(os.path.basename(glob_pattern[i]))[-1] in ['.jpg', '.jpeg', '.png']:
            out.append(glob_pattern[i])
    return out


def resizeAndPad(img, size, padColor=0):
    h, w = img.shape[:2]
    sh, sw = size

    # interpolation method
    if h > sh or w > sw:  # shrinking image
        interp = cv2.INTER_AREA
    else:  # stretching image
        interp = cv2.INTER_CUBIC

    # aspect ratio of image
    aspect = w / h  # if on Python 2, you might need to cast as a float: float(w)/h

    # compute scaling and pad sizing
    if aspect > 1:  # horizontal image
        new_w = sw
        new_h = np.round(new_w / aspect).astype(int)
        pad_vert = (sh - new_h) / 2
        pad_top, pad_bot = np.floor(pad_vert).astype(int), np.ceil(pad_vert).astype(int)
        pad_left, pad_right = 0, 0
    elif aspect < 1:  # vertical image
        new_h = sh
        new_w = np.round(new_h * aspect).astype(int)
        pad_horz = (sw - new_w) / 2
        pad_left, pad_right = np.floor(pad_horz).astype(int), np.ceil(pad_horz).astype(int)
        pad_top, pad_bot = 0, 0
    else:  # square image
        new_h, new_w = sh, sw
        pad_left, pad_right, pad_top, pad_bot = 0, 0, 0, 0

    # set pad color
    if len(img.shape) == 3 and not isinstance(padColor, (list, tuple, np.ndarray)):  # color image but only one color provided
        padColor = [padColor] * 3

    # scale and pad
    scaled_img = cv2.resize(img, (new_w, new_h), interpolation=interp)
    scaled_img = cv2.copyMakeBorder(scaled_img, pad_top, pad_bot, pad_left, pad_right, borderType=cv2.BORDER_CONSTANT, value=padColor)

    return scaled_img


class Ui_MainWindow(object):
    def setupUi(self, MainWindow):
        self.folder_contents = None
        self.folder_contents_i = None

        MainWindow.setObjectName("MainWindow")
        MainWindow.setFixedSize(632, 447)
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Fixed, QtWidgets.QSizePolicy.Fixed)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(0)
        sizePolicy.setHeightForWidth(MainWindow.sizePolicy().hasHeightForWidth())
        MainWindow.setSizePolicy(sizePolicy)
        MainWindow.setAnimated(False)
        MainWindow.setDocumentMode(False)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
        # self.graphicsView = QtWidgets.QGraphicsView(self.centralwidget)
        # self.graphicsView.setGeometry(QtCore.QRect(10, 11, 611, 311))
        # self.graphicsView.setObjectName("graphicsView")

        self.lineEdit = QtWidgets.QLineEdit(self.centralwidget)
        self.lineEdit.setGeometry(QtCore.QRect(10, 330, 611, 31))
        self.lineEdit.setObjectName("lineEdit")
        self.lineEdit.returnPressed.connect(self.write_txtbox)
        self.lineEdit.setEnabled(False)

        self.comboBox = QtWidgets.QComboBox(self.centralwidget)
        self.comboBox.setGeometry(QtCore.QRect(10, 370, 181, 25))
        self.comboBox.setObjectName("comboBox")
        self.comboBox.addItem("")
        self.comboBox.addItem("")
        self.comboBox.setEnabled(False)

        self.save_button = QtWidgets.QPushButton(self.centralwidget)
        self.save_button.setGeometry(QtCore.QRect(530, 370, 89, 25))
        self.save_button.setObjectName("save_button")
        self.save_button.clicked.connect(self.write_txtbox)
        self.save_button.setEnabled(False)

        self.next_button = QtWidgets.QPushButton(self.centralwidget)
        self.next_button.setGeometry(QtCore.QRect(400, 370, 31, 25))
        self.next_button.setObjectName("next_button")
        self.next_button.clicked.connect(self.next_image)
        self.next_button.setEnabled(False)

        self.last_button = QtWidgets.QPushButton(self.centralwidget)
        self.last_button.setGeometry(QtCore.QRect(360, 370, 31, 25))
        self.last_button.setObjectName("last_button")
        self.last_button.clicked.connect(self.prev_image)
        self.last_button.setEnabled(False)

        self.delete_button = QtWidgets.QPushButton(self.centralwidget)
        self.delete_button.setGeometry(QtCore.QRect(440, 370, 89, 25))
        self.delete_button.setObjectName("delete_button")
        self.delete_button.clicked.connect(self.delete_image)
        self.delete_button.setEnabled(False)

        MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 632, 22))
        self.menubar.setObjectName("menubar")
        self.menuFile = QtWidgets.QMenu(self.menubar)
        self.actionQuit = QtWidgets.QAction(MainWindow)
        MainWindow.setMenuBar(self.menubar)
        self.statusBar = QtWidgets.QStatusBar(MainWindow)
        self.statusBar.setSizeGripEnabled(False)
        self.statusBar.setObjectName("statusBar")
        MainWindow.setStatusBar(self.statusBar)
        self.actionQuit.setObjectName("actionQuit")
        self.actionQuit.triggered.connect(self.close)

        self.actionSave = QtWidgets.QAction(MainWindow)
        self.actionSave.setObjectName("actionSave")
        self.actionSave.triggered.connect(self.write_txtbox)

        self.actionNext = QtWidgets.QAction(MainWindow)
        self.actionNext.setObjectName("actionNext")
        self.actionNext.triggered.connect(self.prev_image)

        self.actionOpen = QtWidgets.QAction(MainWindow)
        self.actionOpen.setObjectName("actionOpen")
        self.actionOpen.triggered.connect(self.folder_open)

        self.actionNext = QtWidgets.QAction(MainWindow)
        self.actionNext.setObjectName("actionNext")
        self.actionNext.triggered.connect(self.next_image)

        self.actionDelete = QtWidgets.QAction(MainWindow)
        self.actionDelete.setObjectName("actionDelete")
        self.actionDelete.triggered.connect(self.delete_image)

        self.menuFile.addAction(self.actionOpen)
        self.menuFile.addAction(self.actionQuit)
        self.menuFile.addSeparator()
        self.menuFile.addAction(self.actionNext)
        self.menuFile.addAction(self.actionNext)
        self.menuFile.addAction(self.actionDelete)
        self.menuFile.addAction(self.actionSave)
        self.menubar.addAction(self.menuFile.menuAction())

        self.label = QLabel(self)
        self.label.setGeometry(QtCore.QRect(10, 30, 611, 311))

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "Bulk Image Tagger"))
        self.lineEdit.setPlaceholderText(_translate("MainWindow", "Enter prompt here"))
        self.comboBox.setItemText(0, _translate("MainWindow", "Premade Prompts"))
        self.comboBox.setItemText(1, _translate("MainWindow", "example1"))
        self.save_button.setText(_translate("MainWindow", "Save"))
        self.next_button.setText(_translate("MainWindow", "→"))
        self.last_button.setText(_translate("MainWindow", "←"))
        self.delete_button.setText(_translate("MainWindow", "Delete"))
        self.menuFile.setTitle(_translate("MainWindow", "File"))
        self.actionQuit.setText(_translate("MainWindow", "Quit"))
        self.actionQuit.setShortcut(_translate("MainWindow", "Ctrl+Q"))
        self.actionSave.setText(_translate("MainWindow", "&Save"))
        self.actionSave.setShortcut(_translate("MainWindow", "Ctrl+S"))
        self.actionNext.setText(_translate("MainWindow", "Previous"))
        self.actionNext.setShortcut(_translate("MainWindow", "Right"))
        self.actionOpen.setText(_translate("MainWindow", "Open Folder"))
        self.actionOpen.setShortcut(_translate("MainWindow", "Ctrl+O"))
        self.actionDelete.setText(_translate("MainWindow", "&Delete"))
        self.actionDelete.setShortcut(_translate("MainWindow", "Ctrl+D"))

    def folder_open(self):
        self.user_dir = QFileDialog.getExistingDirectory(
            self, "Open Directory",
            "",
            QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks,
        )
        self.folder_contents = get_folder(self.user_dir)
        self.folder_contents_i = 0
        self.last_button.setEnabled(False)
        if len(self.folder_contents) > 0:
            self.load_image(self.folder_contents[self.folder_contents_i])
            self.enable(True)
        else:
            self.statusBar.showMessage('No images found!')

    def load_image(self, image_path):
        self.lineEdit.clear()
        if os.path.isfile(image_path):
            im = cv2.imread(image_path)  # square image
            im_resize = resizeAndPad(im, (311, 311), 311)
            is_success, im_buf_arr = cv2.imencode(".png", im_resize)
            im = Image.open(BytesIO(im_buf_arr.tobytes()))
            im.tobytes('raw', 'RGB')
            qim = ImageQt(im)
            pixmap = QPixmap.fromImage(qim)
            self.label.setPixmap(pixmap)
            self.label.show()
            self.statusBar.showMessage(os.path.split(self.folder_contents[self.folder_contents_i])[-1])

    def next_image(self):
        if self.folder_contents is not None:
            if self.folder_contents_i + 1 < len(self.folder_contents):
                self.folder_contents_i += 1
                self.load_image(self.folder_contents[self.folder_contents_i])
            else:
                self.next_button.setEnabled(False)
            if self.folder_contents_i - 1 < len(self.folder_contents):
                self.last_button.setEnabled(True)

    def prev_image(self):
        if self.folder_contents is not None:
            if self.folder_contents_i - 1 >= 0:
                self.folder_contents_i -= 1
                self.load_image(self.folder_contents[self.folder_contents_i])
            elif self.folder_contents_i - 1 < len(self.folder_contents):
                self.last_button.setEnabled(False)
            if self.folder_contents_i + 1 < len(self.folder_contents):
                self.next_button.setEnabled(True)

    def write_txtbox(self):
        if self.folder_contents is not None:
            textbox_content = self.lineEdit.text()
            if len(textbox_content) > 0:
                f = open(f'{self.folder_contents[self.folder_contents_i]}.txt', 'w')
                f.write(textbox_content)
                f.close()
                self.lineEdit.clear()
                self.next_image()

    def delete_image(self):
        if self.folder_contents is not None and len(self.folder_contents) > 0:
            self.label.hide()
            try:
                os.remove(self.folder_contents[self.folder_contents_i])
            except:
                print(self.folder_contents, self.folder_contents_i)
            if os.path.exists(f'{self.folder_contents[self.folder_contents_i]}.txt'):
                os.remove(f'{self.folder_contents[self.folder_contents_i]}.txt')
            self.statusBar.showMessage(f'Deleted {os.path.split(self.folder_contents[self.folder_contents_i])[-1]}')
            del self.folder_contents[self.folder_contents_i]
            self.folder_contents_i -= 1
            if len(self.folder_contents) > 0:
                self.next_image()
            else:
                self.enable(False)

    def enable(self, action):
        self.save_button.setEnabled(action)
        self.next_button.setEnabled(action)
        self.last_button.setEnabled(action)
        self.delete_button.setEnabled(action)
        self.lineEdit.setEnabled(action)
        self.comboBox.setEnabled(action)