class AddressWidget(QTabWidget):
""" The central widget of the application. Most of the addressbook's
functionality is contained in this class.
"""
selectionChanged = Signal(QItemSelection)
def __init__(self, parent=None):
""" Initialize the AddressWidget. """
super(AddressWidget, self).__init__(parent)
self.tableModel = TableModel()
self.newAddressTab = NewAddressTab()
self.newAddressTab.sendDetails.connect(self.addEntry)
self.addTab(self.newAddressTab, "Address Book")
self.setupTabs()
def addEntry(self, name=None, address=None):
""" Add an entry to the addressbook. """
if name is None and address is None:
addDialog = AddDialogWidget()
if addDialog.exec_():
name = addDialog.name
address = addDialog.address
address = {"name": name, "address": address}
addresses = self.tableModel.addresses[:]
# The QT docs for this example state that what we're doing here
# is checking if the entered name already exists. What they
# (and we here) are actually doing is checking if the whole
# name/address pair exists already - ok for the purposes of this
# example, but obviously not how a real addressbook application
# should behave.
try:
addresses.remove(address)
QMessageBox.information(self, "Duplicate Name",
"The name \"%s\" already exists." % name)
except ValueError:
# The address didn't already exist, so let's add it to the model.
# Step 1: create the row
self.tableModel.insertRows(0)
# Step 2: get the index of the newly created row and use it.
# to set the name
ix = self.tableModel.index(0, 0, QModelIndex())
self.tableModel.setData(ix, address["name"], Qt.EditRole)
# Step 3: lather, rinse, repeat for the address.
ix = self.tableModel.index(0, 1, QModelIndex())
self.tableModel.setData(ix, address["address"], Qt.EditRole)
# Remove the newAddressTab, as we now have at least one
# address in the model.
self.removeTab(self.indexOf(self.newAddressTab))
# The screenshot for the QT example shows nicely formatted
# multiline cells, but the actual application doesn't behave
# quite so nicely, at least on Ubuntu. Here we resize the newly
# created row so that multiline addresses look reasonable.
tableView = self.currentWidget()
tableView.resizeRowToContents(ix.row())
def editEntry(self):
""" Edit an entry in the addressbook. """
tableView = self.currentWidget()
proxyModel = tableView.model()
selectionModel = tableView.selectionModel()
# Get the name and address of the currently selected row.
indexes = selectionModel.selectedRows()
for index in indexes:
row = proxyModel.mapToSource(index).row()
ix = self.tableModel.index(row, 0, QModelIndex())
name = self.tableModel.data(ix, Qt.DisplayRole)
ix = self.tableModel.index(row, 1, QModelIndex())
address = self.tableModel.data(ix, Qt.DisplayRole)
# Open an addDialogWidget, and only allow the user to edit the address.
addDialog = AddDialogWidget()
addDialog.setWindowTitle("Edit a Contact")
addDialog.nameText.setReadOnly(True)
addDialog.nameText.setText(name)
addDialog.addressText.setText(address)
# If the address is different, add it to the model.
if addDialog.exec_():
newAddress = addDialog.address
if newAddress != address:
ix = self.tableModel.index(row, 1, QModelIndex())
self.tableModel.setData(ix, newAddress, Qt.EditRole)
def removeEntry(self):
""" Remove an entry from the addressbook. """
tableView = self.currentWidget()
proxyModel = tableView.model()
selectionModel = tableView.selectionModel()
# Just like editEntry, but this time remove the selected row.
indexes = selectionModel.selectedRows()
for index in indexes:
row = proxyModel.mapToSource(index).row()
self.tableModel.removeRows(row)
# If we've removed the last address in the model, display the
# newAddressTab
if self.tableModel.rowCount() == 0:
self.insertTab(0, self.newAddressTab, "Address Book")
def setupTabs(self):
""" Setup the various tabs in the AddressWidget. """
groups = ["ABC", "DEF", "GHI", "JKL", "MNO", "PQR", "STU", "VW", "XYZ"]
for group in groups:
proxyModel = QSortFilterProxyModel(self)
proxyModel.setSourceModel(self.tableModel)
proxyModel.setDynamicSortFilter(True)
tableView = QTableView()
tableView.setModel(proxyModel)
tableView.setSortingEnabled(True)
tableView.setSelectionBehavior(QAbstractItemView.SelectRows)
tableView.horizontalHeader().setStretchLastSection(True)
tableView.verticalHeader().hide()
tableView.setEditTriggers(QAbstractItemView.NoEditTriggers)
tableView.setSelectionMode(QAbstractItemView.SingleSelection)
# This here be the magic: we use the group name (e.g. "ABC") to
# build the regex for the QSortFilterProxyModel for the group's
# tab. The regex will end up looking like "^[ABC].*", only
# allowing this tab to display items where the name starts with
# "A", "B", or "C". Notice that we set it to be case-insensitive.
reFilter = "^[%s].*" % group
proxyModel.setFilterRegExp(QRegExp(reFilter, Qt.CaseInsensitive))
proxyModel.setFilterKeyColumn(0) # Filter on the "name" column
proxyModel.sort(0, Qt.AscendingOrder)
# This prevents an application crash (see: http://www.qtcentre.org/threads/58874-QListView-SelectionModel-selectionChanged-Crash)
viewselectionmodel = tableView.selectionModel()
tableView.selectionModel().selectionChanged.connect(self.selectionChanged)
self.addTab(tableView, group)
# Note: the QT example uses a QDataStream for the saving and loading.
# Here we're using a python dictionary to store the addresses, which
# can't be streamed using QDataStream, so we just use cpickle for this
# example.
def readFromFile(self, filename):
""" Read contacts in from a file. """
try:
f = open(filename, "rb")
addresses = pickle.load(f)
except IOError:
QMessageBox.information(self, "Unable to open file: %s" % filename)
finally:
f.close()
if len(addresses) == 0:
QMessageBox.information(self, "No contacts in file: %s" % filename)
else:
for address in addresses:
self.addEntry(address["name"], address["address"])
def writeToFile(self, filename):
""" Save all contacts in the model to a file. """
try:
f = open(filename, "wb")
pickle.dump(self.tableModel.addresses, f)
except IOError:
QMessageBox.information(self, "Unable to open file: %s" % filename)
finally:
f.close()
class AddressWidget(QTabWidget):
""" The central widget of the application. Most of the addressbook's
functionality is contained in this class.
"""
selectionChanged = Signal(QItemSelection)
def __init__(self, parent=None):
""" Initialize the AddressWidget. """
super(AddressWidget, self).__init__(parent)
self.tableModel = TableModel()
self.newAddressTab = NewAddressTab()
self.newAddressTab.sendDetails.connect(self.addEntry)
self.addTab(self.newAddressTab, "Address Book")
self.setupTabs()
def addEntry(self, name=None, address=None):
""" Add an entry to the addressbook. """
if name is None and address is None:
addDialog = AddDialogWidget()
if addDialog.exec_():
name = addDialog.name
address = addDialog.address
address = {"name": name, "address": address}
addresses = self.tableModel.addresses[:]
# The QT docs for this example state that what we're doing here
# is checking if the entered name already exists. What they
# (and we here) are actually doing is checking if the whole
# name/address pair exists already - ok for the purposes of this
# example, but obviously not how a real addressbook application
# should behave.
try:
addresses.remove(address)
QMessageBox.information(self, "Duplicate Name",
"The name \"%s\" already exists." % name)
except ValueError:
# The address didn't already exist, so let's add it to the model.
# Step 1: create the row
self.tableModel.insertRows(0)
# Step 2: get the index of the newly created row and use it.
# to set the name
ix = self.tableModel.index(0, 0, QModelIndex())
self.tableModel.setData(ix, address["name"], Qt.EditRole)
# Step 3: lather, rinse, repeat for the address.
ix = self.tableModel.index(0, 1, QModelIndex())
self.tableModel.setData(ix, address["address"], Qt.EditRole)
# Remove the newAddressTab, as we now have at least one
# address in the model.
self.removeTab(self.indexOf(self.newAddressTab))
# The screenshot for the QT example shows nicely formatted
# multiline cells, but the actual application doesn't behave
# quite so nicely, at least on Ubuntu. Here we resize the newly
# created row so that multiline addresses look reasonable.
tableView = self.currentWidget()
tableView.resizeRowToContents(ix.row())
def editEntry(self):
""" Edit an entry in the addressbook. """
tableView = self.currentWidget()
proxyModel = tableView.model()
selectionModel = tableView.selectionModel()
# Get the name and address of the currently selected row.
indexes = selectionModel.selectedRows()
for index in indexes:
row = proxyModel.mapToSource(index).row()
ix = self.tableModel.index(row, 0, QModelIndex())
name = self.tableModel.data(ix, Qt.DisplayRole)
ix = self.tableModel.index(row, 1, QModelIndex())
address = self.tableModel.data(ix, Qt.DisplayRole)
# Open an addDialogWidget, and only allow the user to edit the address.
addDialog = AddDialogWidget()
addDialog.setWindowTitle("Edit a Contact")
addDialog.nameText.setReadOnly(True)
addDialog.nameText.setText(name)
addDialog.addressText.setText(address)
# If the address is different, add it to the model.
if addDialog.exec_():
newAddress = addDialog.address
if newAddress != address:
ix = self.tableModel.index(row, 1, QModelIndex())
self.tableModel.setData(ix, newAddress, Qt.EditRole)
def removeEntry(self):
""" Remove an entry from the addressbook. """
tableView = self.currentWidget()
proxyModel = tableView.model()
selectionModel = tableView.selectionModel()
# Just like editEntry, but this time remove the selected row.
indexes = selectionModel.selectedRows()
for index in indexes:
row = proxyModel.mapToSource(index).row()
self.tableModel.removeRows(row)
# If we've removed the last address in the model, display the
# newAddressTab
if self.tableModel.rowCount() == 0:
self.insertTab(0, self.newAddressTab, "Address Book")
def setupTabs(self):
""" Setup the various tabs in the AddressWidget. """
groups = ["ABC", "DEF", "GHI", "JKL", "MNO", "PQR", "STU", "VW", "XYZ"]
for group in groups:
proxyModel = QSortFilterProxyModel(self)
proxyModel.setSourceModel(self.tableModel)
proxyModel.setDynamicSortFilter(True)
tableView = QTableView()
tableView.setModel(proxyModel)
tableView.setSortingEnabled(True)
tableView.setSelectionBehavior(QAbstractItemView.SelectRows)
tableView.horizontalHeader().setStretchLastSection(True)
tableView.verticalHeader().hide()
tableView.setEditTriggers(QAbstractItemView.NoEditTriggers)
tableView.setSelectionMode(QAbstractItemView.SingleSelection)
# This here be the magic: we use the group name (e.g. "ABC") to
# build the regex for the QSortFilterProxyModel for the group's
# tab. The regex will end up looking like "^[ABC].*", only
# allowing this tab to display items where the name starts with
# "A", "B", or "C". Notice that we set it to be case-insensitive.
reFilter = "^[%s].*" % group
proxyModel.setFilterRegExp(QRegExp(reFilter, Qt.CaseInsensitive))
proxyModel.setFilterKeyColumn(0) # Filter on the "name" column
proxyModel.sort(0, Qt.AscendingOrder)
#tableView.selectionModel().selectionChanged.connect(self.selectionChanged)
self.addTab(tableView, group)
# Note: the QT example uses a QDataStream for the saving and loading.
# Here we're using a python dictionary to store the addresses, which
# can't be streamed using QDataStream, so we just use cpickle for this
# example.
def readFromFile(self, filename):
""" Read contacts in from a file. """
try:
f = open(filename, "rb")
addresses = pickle.load(f)
except IOError:
QMessageBox.information(self, "Unable to open file: %s" % filename)
finally:
f.close()
if len(addresses) == 0:
QMessageBox.information(self, "No contacts in file: %s" % filename)
else:
for address in addresses:
self.addEntry(address["name"], address["address"])
def writeToFile(self, filename):
""" Save all contacts in the model to a file. """
try:
f = open(filename, "wb")
pickle.dump(self.tableModel.addresses, f)
except IOError:
QMessageBox.information(self, "Unable to open file: %s" % filename)
finally:
f.close()
class ImageSimilarity(QWidget):
"""
Main class
"""
def __init__(self, parent=None):
super(ImageSimilarity, self).__init__(parent)
#Build window
self.setWindowTitle('Image Similarity Search')
self.radio_pca = QRadioButton('PCA')
self.radio_pca.setChecked(True)
self.radio_deep = QRadioButton('Convolutional')
self.button_load_weights = QPushButton('Load Weights')
self.button_build_DB = QPushButton('Build Database')
self.button_save_DB = QPushButton('Save Database')
self.button_load_DB = QPushButton('Load Database')
self.button_load_img = QPushButton('Load Image')
self.button_search = QPushButton('Search')
self.button_search.setEnabled(False)
self.button_save_result = QPushButton('Save Result')
self.button_save_result.setEnabled(False)
self.spinbox = QSpinBox()
self.spinbox.setMinimum(-1)
self.spinbox.setMaximum(99999)
self.spinbox.setValue(0)
self.spinboxlabel = QLabel('Max. Distance:')
self.spinboxlabel.setAlignment(Qt.AlignRight)
self.searchlabel = QLabel('Query Image:')
self.countlabel = QLabel('# of Hits:')
self.table = QTableView()
#self.table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
self.table.horizontalHeader().setStretchLastSection(True)
self.model = TableModel(parent=self)
self.table.setModel(self.model)
self.table.setSortingEnabled(True)
self.table.setEditTriggers(QAbstractItemView.NoEditTriggers)
hbox_radio = QHBoxLayout()
hbox_radio.addWidget(self.radio_pca)
hbox_radio.addWidget(self.radio_deep)
hbox_radio.addWidget(self.button_load_weights)
hbox_buttons = QHBoxLayout()
hbox_buttons.addWidget(self.button_build_DB)
hbox_buttons.addWidget(self.button_load_DB)
hbox_buttons.addWidget(self.button_save_DB)
hbox_buttons.addWidget(self.button_load_img)
hbox_buttons.addWidget(self.button_search)
hbox_buttons.addWidget(self.spinboxlabel)
hbox_buttons.addWidget(self.spinbox)
radio_widget = QWidget()
radio_widget.setLayout(hbox_radio)
button_widget = QWidget()
button_widget.setLayout(hbox_buttons)
external_layout = QVBoxLayout()
external_layout.addWidget(radio_widget)
external_layout.addWidget(button_widget)
external_layout.addWidget(self.searchlabel)
external_layout.addWidget(self.countlabel)
external_layout.addWidget(self.table)
external_layout.addWidget(self.button_save_result)
self.setLayout(external_layout)
#Connect signals and slots
self.radio_pca.toggled.connect(lambda: self.switch_NN(self.radio_pca))
self.radio_deep.toggled.connect(
lambda: self.switch_NN(self.radio_deep))
self.button_load_weights.clicked.connect(self.load_weights)
self.button_build_DB.clicked.connect(self.build_DB)
self.button_load_DB.clicked.connect(self.load_DB)
self.button_save_DB.clicked.connect(self.save_DB)
self.button_load_img.clicked.connect(self.load_image)
self.button_search.clicked.connect(self.search)
self.button_save_result.clicked.connect(self.save_search)
#Build neural network
self.encoder = build_pca_autoencoder()
#See if a database already exists
if not os.path.isfile(os.path.join(application_path, 'std_db_pca.db')):
self.firstTimeStart(os.path.join(application_path,
'std_db_pca.db'))
else:
self.database = load_database(
os.path.join(application_path, 'std_db_pca.db'))
def firstTimeStart(self, savename):
msg = QMessageBox()
msg.setIcon(QMessageBox.Information)
msg.setText(
"Please select a directory to search for images in. This will then build the search database."
)
msg.setWindowTitle("Setup")
msg.setStandardButtons(QMessageBox.Ok)
if msg.exec_():
dir = None
while not dir:
dir = QFileDialog.getExistingDirectory(
self, 'Choose Search Root', QDir.homePath(),
QFileDialog.ShowDirsOnly)
paths = get_image_paths(dir)
self.progresswindow = self.make_progress_window(len(paths))
self.progresswindow.show()
graph = tf.get_default_graph()
build_thread = Build_Thread(self.encoder, paths, graph, self)
build_thread.processed.connect(self.update_progress)
build_thread.finished.connect(self.set_database)
build_thread.finished.connect(
lambda db: save_database(db, savename))
build_thread.close.connect(self.progresswindow.close)
build_thread.close.connect(lambda: self.build_message(savename))
build_thread.start()
def build_message(self, savename=None):
msg = QMessageBox()
msg.setIcon(QMessageBox.Information)
msg.setText("Database complete!")
if not savename:
msg.setInformativeText(
"Don't forget to save your database to avoid having to rebuild it upon program restart."
)
else:
msg.setInformativeText(
"The database was saved as {} and will be loaded as default on all subsequent program starts."
.format(savename))
msg.setWindowTitle("Completed")
msg.setStandardButtons(QMessageBox.Ok)
msg.exec_()
def make_progress_window(self, max):
window = QDialog(self)
window.setWindowTitle("Building Search Database")
self.progress = QProgressBar(window)
self.progress.setMaximum(max)
self.progresslabel = QLabel(window)
self.progresslabel.setText("Processed: 0 of {}".format(max))
layout = QVBoxLayout()
window.setLayout(layout)
layout.addWidget(self.progresslabel)
layout.addWidget(self.progress)
window.setModal(Qt.ApplicationModal)
return window
def update_progress(self, prog, max):
self.progresslabel.setText("Processed: {} of {}".format(prog, max))
self.progress.setValue(prog)
def set_database(self, db):
self.database = db
def load_image(self):
f, _ = QFileDialog.getOpenFileName(
self, 'Load Image', QDir.homePath(),
"Image Files (*.jpg *.jpeg *.png *.tga *.bmp)")
if f:
self.search_img_path = f
self.searchlabel.setText('Query Image: {}'.format(
os.path.basename(f)))
if not self.button_search.isEnabled():
self.button_search.setEnabled(True)
def search(self):
if self.search_img_path:
#start = timer()
self.lastSearchQuery = self.search_img_path
if not self.button_save_result.isEnabled():
self.button_save_result.setEnabled(True)
hits = similarity_search(self.search_img_path, self.database,
self.encoder, self.spinbox.value())
self.countlabel.setText('# of hits: {}'.format(str(len(hits))))
#Update Table Model
self.model.removeRows(0, self.model.rowCount())
for hit in hits:
if b'name' in hit.keys():
row = {
"name": hit[b'name'],
"path": hit[b'path'],
"distance": hit['distance']
}
else:
row = {
"name": hit['name'],
"path": hit['path'],
"distance": hit['distance']
}
self.model.appendRow(row)
#end = timer()
#print('Search Time: {} seconds'.format(end - start))
def save_search(self):
filename, _ = QFileDialog.getSaveFileName(self, 'Save Result',
QDir.currentPath(),
'Text files (*.txt)')
if filename:
with open(filename, "w") as file:
file.write("Query image: {}\n\n".format(self.lastSearchQuery))
for i in range(self.model.rowCount()):
file.write("Name: {}, Path: {}, Distance: {}\n".format(
self.model.results[i]['name'],
self.model.results[i]['path'],
self.model.results[i]['distance']))
def switch_NN(self, b):
if b.text() == 'PCA' and b.isChecked():
reset_tf_session()
self.encoder = build_pca_autoencoder()
self.database = load_database(
os.path.join(application_path, 'std_db_pca.db'))
self.search_img_path = None
self.searchlabel.setText('Query Image:')
self.button_search.setEnabled(False)
elif b.text() == 'Convolutional' and b.isChecked():
reset_tf_session()
self.encoder = build_deep_autoencoder()
if os.path.isfile(os.path.join(application_path,
'std_db_deep.db')):
self.database = load_database(
os.path.join(application_path, 'std_db_deep.db'))
else:
self.firstTimeStart(
os.path.join(application_path, 'std_db_deep.db'))
self.search_img_path = None
self.searchlabel.setText('Query Image:')
self.button_search.setEnabled(False)
else:
pass
def load_weights(self):
filename, _ = QFileDialog.getOpenFileName(
self, 'Load Weights', QDir.currentPath(),
'Neural Network Weights (*.h5)')
if filename:
load_NN_weights(self.encoder, filename)
def build_DB(self):
dlg = QFileDialog()
dlg.setFileMode(QFileDialog.Directory)
if dlg.exec_():
#start = timer()
dirs = dlg.selectedFiles()
paths = get_image_paths(dirs[0])
self.progresswindow = self.make_progress_window(len(paths))
self.progresswindow.show()
graph = tf.get_default_graph()
build_thread = Build_Thread(self.encoder, paths, graph, self)
build_thread.processed.connect(self.update_progress)
build_thread.finished.connect(self.set_database)
build_thread.close.connect(self.progresswindow.close)
build_thread.close.connect(self.build_message)
#build_thread.close.connect(lambda : print(timer() - start))
build_thread.start()
def load_DB(self):
dbname, _ = QFileDialog.getOpenFileName(self, 'Load Database',
QDir.currentPath(),
"Database files (*.db)")
if dbname:
self.database = load_database(dbname)
def save_DB(self):
filename, _ = QFileDialog.getSaveFileName(self, 'Save Database',
QDir.currentPath(),
'Database files (*.db)')
if filename:
save_database(self.database, filename, True)