diff --git a/AppGuiFunctions.py b/AppGuiFunctions.py index a207729..cbd70f7 100644 --- a/AppGuiFunctions.py +++ b/AppGuiFunctions.py @@ -1,3 +1,5 @@ +import os + import cv2 import numpy as np import shutil @@ -27,37 +29,44 @@ def saveFile(self, savepath): def imageFitter(self, file_input_path, file_output_path, winwidth, winheight): rawimg = cv2.imread(file_input_path, 1) + if rawimg is None: + raise FileNotFoundError(f"Could not read image: {file_input_path}") self.imgrealwidth = rawimg.shape[1] self.imgrealheight = rawimg.shape[0] - ##Will change the logic below later, for now its been configured for square window shapes :P + # Will change the logic below later, for now it's been configured for square window shapes - if rawimg.shape[1] >= rawimg.shape[0] : + if rawimg.shape[1] >= rawimg.shape[0]: newx = winwidth - newy = (rawimg.shape[0]/rawimg.shape[1])*winwidth + newy = (rawimg.shape[0] / rawimg.shape[1]) * winwidth else: newy = winheight - newx = (rawimg.shape[1]/rawimg.shape[0])*winheight - + newx = (rawimg.shape[1] / rawimg.shape[0]) * winheight + newdim = (int(newx), int(newy)) - resizedimg = cv2.resize(rawimg, newdim, interpolation = cv2.INTER_AREA) + resizedimg = cv2.resize(rawimg, newdim, interpolation=cv2.INTER_AREA) + # Ensure output directory exists + os.makedirs(os.path.dirname(os.path.abspath(file_output_path)), exist_ok=True) cv2.imwrite(file_output_path, resizedimg) def superresUpscaler(self, model): self.finaloutputpath = "processimage/finalprocessedimage_" + model + ".jpg" self.displayoutput = "processimage/downscaled_finalprocessedimage_" + model + ".jpg" - + self.ganPredict = DeepLearningGANModels(model, self.path_to_file) model_img = self.ganPredict.modelBash() - array_img = np.array(model_img).astype(np.float32) + if model_img is None: + return False + + # Ensure output directory exists + os.makedirs(os.path.dirname(os.path.abspath(self.finaloutputpath)), exist_ok=True) + + array_img = np.array(model_img).astype(np.uint8) processed_model_img = cv2.cvtColor(array_img, cv2.COLOR_RGB2BGR) cv2.imwrite(self.finaloutputpath, processed_model_img) self.imageFitter(self.finaloutputpath, self.displayoutput, self.winwidth, self.winheight) - if model_img is not None : - return True - else : - return False \ No newline at end of file + return True \ No newline at end of file diff --git a/AppWindowQt.py b/AppWindowQt.py index 3929d6c..10011c0 100644 --- a/AppWindowQt.py +++ b/AppWindowQt.py @@ -1,8 +1,10 @@ # -*- coding: utf-8 -*- +import os +import sys + from PyQt5 import QtCore, QtGui, QtWidgets from PyQt5.QtWidgets import QFileDialog, QMessageBox, QLabel, QMainWindow from AppGuiFunctions import GuiFunctions -import sys class QLabel_alterada(QLabel): @@ -167,14 +169,21 @@ def retranslateUi(self, MainWindow): self.actionSave.setShortcut(_translate("MainWindow", "Ctrl+S")) def processImage(self, model): - if not self.file_has_been_loaded : + if not self.file_has_been_loaded: self.showNoFilePopup() - else : - successRes = self.gui_functions.superresUpscaler(model) - if successRes : - self.processed_photo.setPixmap(QtGui.QPixmap("processimage/downscaled_finalprocessedimage_" + model + ".jpg")) + else: + try: + successRes = self.gui_functions.superresUpscaler(model) + except Exception: + self.showFatalPopup() + return + + if successRes: + self.processed_photo.setPixmap(QtGui.QPixmap( + "processimage/downscaled_finalprocessedimage_" + model + ".jpg" + )) self.photo_has_been_processed = True - else : + else: self.showFatalPopup() def displayChosenImage(self): @@ -189,32 +198,38 @@ def openExtraWindow(self): def openFileDialogBox(self): file_path = QFileDialog.getOpenFileName()[0] - extension = file_path.split('.')[-1] ##Use this for checking filetype and throwing errors and stuff - - if extension in self.accepted_extensions : + if not file_path: + return # user cancelled + + _, ext = os.path.splitext(file_path) + extension = ext.lstrip(".").lower() + + if extension in self.accepted_extensions: self.gui_functions.openFile(file_path, self.imgwidth, self.imgheight) self.file_has_been_loaded = True self.displayChosenImage() self.processed_photo.setText(QtCore.QCoreApplication.translate("MainWindow", "Click Here to View Your UpScaled Image")) self.photo_has_been_processed = False - elif extension == "" : - pass - else : + else: self.showExtensionsMismatchPopup(extension) def saveFileDialogBox(self): - if not self.file_has_been_loaded : + if not self.file_has_been_loaded: self.showNoFilePopup() - else : - file_path = QFileDialog.getSaveFileName()[0] - extension = file_path.split('.')[-1] ##Use this for checking filetype and throwing errors and stuff - - if extension == file_path : - file_path = file_path + ".jpg" - self.gui_functions.saveFile(file_path) - else : - self.gui_functions.saveFile(file_path) + return + if not self.photo_has_been_processed: + self.showNoUpscalingDonePopup() + return + + file_path = QFileDialog.getSaveFileName()[0] + if not file_path: + return # user cancelled + + if not os.path.splitext(file_path)[1]: + file_path = file_path + ".jpg" + + self.gui_functions.saveFile(file_path) def showNoFilePopup(self): popmsg = QMessageBox() diff --git a/EnhancerGANModel.py b/EnhancerGANModel.py index 7ba8974..fe2d7ab 100644 --- a/EnhancerGANModel.py +++ b/EnhancerGANModel.py @@ -16,7 +16,8 @@ def __init__(self, model, input_path): self.input_path = input_path def load_image(self, path): - return np.array(Image.open(path))[:,:,:3] + img = Image.open(path).convert("RGB") + return np.array(img) def modelBash(self): if self.modelname == "srganx4" : @@ -27,23 +28,31 @@ def modelBash(self): return None def wdsr_b_finetuned(self): - weights_dir = 'model/weights/wdsr' + weights_dir = os.path.join(os.path.dirname(__file__), 'model', 'weights', 'wdsr') + weights_path = os.path.join(weights_dir, 'weights-wdsr-b-fine-tuned.h5') + if not os.path.isfile(weights_path): + raise FileNotFoundError( + f"WDSR weight file not found: {weights_path}\n" + "Download weights from https://drive.google.com/drive/folders/..." + ) + wdsr_fine_tuned = wdsr_b(scale=4, num_res_blocks=32) - wdsr_fine_tuned.load_weights(os.path.join(weights_dir, 'weights-wdsr-b-fine-tuned.h5')) + wdsr_fine_tuned.load_weights(weights_path) lr = self.load_image(self.input_path) - gan_wdsr_finetuned = resolve_single(wdsr_fine_tuned, lr) - - return gan_wdsr_finetuned + return resolve_single(wdsr_fine_tuned, lr) def srganx4Process(self): - weights_dir = 'model/weights/srgan' - weights_file = lambda filename: os.path.join(weights_dir, filename) + weights_dir = os.path.join(os.path.dirname(__file__), 'model', 'weights', 'srgan') + weights_path = os.path.join(weights_dir, 'gan_generator.h5') + if not os.path.isfile(weights_path): + raise FileNotFoundError( + f"SRGAN weight file not found: {weights_path}\n" + "Download weights from https://drive.google.com/drive/folders/..." + ) gan_generator = srgangenerator() - gan_generator.load_weights(weights_file('gan_generator.h5')) + gan_generator.load_weights(weights_path) lr = self.load_image(self.input_path) - gan_sr = resolve_single(gan_generator, lr) - - return gan_sr \ No newline at end of file + return resolve_single(gan_generator, lr) \ No newline at end of file diff --git a/model/WDSRFineTuned.py b/model/WDSRFineTuned.py index 4b354f4..f7d7c1a 100644 --- a/model/WDSRFineTuned.py +++ b/model/WDSRFineTuned.py @@ -1,10 +1,13 @@ -import tensorflow_addons as tfa - -from tensorflow.python.keras.layers import Add, Conv2D, Input, Lambda -from tensorflow.python.keras.models import Model +from tensorflow.keras.layers import Add, Conv2D, Input, Lambda +from tensorflow.keras.models import Model from model.common import normalize, denormalize, pixel_shuffle +try: + import tensorflow_addons as tfa +except ImportError: + tfa = None + def wdsr_a(scale, num_filters=32, num_res_blocks=8, res_block_expansion=4, res_block_scaling=None): return wdsr(scale, num_filters, num_res_blocks, res_block_expansion, res_block_scaling, res_block_a) @@ -56,4 +59,12 @@ def res_block_b(x_in, num_filters, expansion, kernel_size, scaling): def conv2d_weightnorm(filters, kernel_size, padding='same', activation=None, **kwargs): - return tfa.layers.WeightNormalization(Conv2D(filters, kernel_size, padding=padding, activation=activation, **kwargs), data_init=False) + if tfa is None: + raise ImportError( + "tensorflow-addons is required for WDSR models. " + "Install with: pip install tensorflow-addons" + ) + return tfa.layers.WeightNormalization( + Conv2D(filters, kernel_size, padding=padding, activation=activation, **kwargs), + data_init=False, + ) diff --git a/model/srgan.py b/model/srgan.py index dbc9ecc..91bade3 100644 --- a/model/srgan.py +++ b/model/srgan.py @@ -1,6 +1,6 @@ -from tensorflow.python.keras.layers import Add, BatchNormalization, Conv2D, Dense, Flatten, Input, LeakyReLU, PReLU, Lambda -from tensorflow.python.keras.models import Model -from tensorflow.python.keras.applications.vgg19 import VGG19 +from tensorflow.keras.layers import Add, BatchNormalization, Conv2D, Dense, Flatten, Input, LeakyReLU, PReLU, Lambda +from tensorflow.keras.models import Model +from tensorflow.keras.applications.vgg19 import VGG19 from model.common import pixel_shuffle, normalize_01, normalize_m11, denormalize_m11 diff --git a/run.py b/run.py index a87ea01..4b87a5a 100644 --- a/run.py +++ b/run.py @@ -1,9 +1,25 @@ -from AppWindowQt import Ui_MainWindow -import tensorflow as tf import os +import sys + +from AppWindowQt import Ui_MainWindow + + +def main(): + # Configure GPU memory growth if a GPU is available. + try: + import tensorflow as tf + + physical_devices = tf.config.list_physical_devices("GPU") + if physical_devices: + tf.config.experimental.set_memory_growth(physical_devices[0], True) + except ImportError: + pass # TensorFlow not installed — CPU fallback or inference-only + except Exception: + pass # GPU config failed — continue anyway + + main_ui = Ui_MainWindow() + main_ui.execute() -physical_devices = tf.config.list_physical_devices('GPU') -tf.config.experimental.set_memory_growth(physical_devices[0], True) -main_ui = Ui_MainWindow() -main_ui.execute() \ No newline at end of file +if __name__ == "__main__": + main() \ No newline at end of file