Update app.py

app.py

@@ -6,16 +6,20 @@ import numpy as np
 from PIL import Image
 from model.flol import create_model
 
-
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 #define some auxiliary functions
 pil_to_tensor = transforms.ToTensor()
 
-# Define a list of available weight files
-weights_files = {
-    "flolv2_UHDLL": './weights/flolv2_UHDLL.pt',
-    "flolv2_LOLv2-Real": './weights/flolv2_all_111439.pt',  # Add other available weights here
-}
+# Define a dictionary to map image filenames to weight files
+image_to_weights = {
+    "425_UHD_LL.JPG": './weights/flolv2_UHDLL.pt',
+    "1778_UHD_LL.JPG": './weights/flolv2_UHDLL.pt',  
+    "1791_UHD_LL.JPG": './weights/flolv2_UHDLL.pt',  
+    
+    "low00748.png": './weights/flolv2_all_111439.pt',
+    "low00723.png": './weights/flolv2_all_111439.pt',
+    "low00772.png": './weights/flolv2_all_111439.pt'
+    }
 
 # Initial model setup (without weights)
 model = create_model()
@@ -25,10 +29,10 @@ def load_img(filename):
     img_tensor = pil_to_tensor(img)
     return img_tensor
 
-def process_img(image, weight_file):
-    # Load the selected weight file
-    if weight_file in weights_files:
-        model_path = weights_files[weight_file]
+def process_img(image, filename):
+    # Select the correct weight file based on the image filename
+    if filename in image_to_weights:
+        model_path = image_to_weights[filename]
         checkpoints = torch.load(model_path, map_location=device)
         model.load_state_dict(checkpoints['params'])
         model.to(device)
@@ -59,9 +63,9 @@ Due to the GPU memory limitations, the app might crash if you feed a high-resolu
 '''
 
 examples = [['images/425_UHD_LL.JPG'],
-            ['images/674_UHD_LL.JPG'], 
-            ['images/685_UHD_LL.JPG'], 
-            ['images/945_UHD_LL.JPG'], 
+            ['images/low00772.png'], 
+            ['images/low00723.png'], 
+            ['images/low00748.png'], 
             ['images/1778_UHD_LL.JPG'],
             ['images/1791_UHD_LL.JPG']]
 
@@ -76,8 +80,8 @@ css = """
 demo = gr.Interface(
     fn=process_img,
     inputs=[
-        gr.Image(type='pil', label='input'),
-        gr.Dropdown(choices=list(weights_files.keys()), label='Select Weight File', default="flolv2_UHDLL")
+        gr.Image(type='pil', label='input', tool='editor'),
+        gr.Textbox(label="Image Filename", interactive=False)
     ],
     outputs=[gr.Image(type='pil', label='output')],
     title=title,
@@ -86,5 +90,16 @@ demo = gr.Interface(
     css=css
 )
 
+# Updating the filename in the input after selection
+def update_filename(image):
+    # Retrieve the filename from the input image
+    if image:
+        filename = image.filename  # Gradio automatically gives the file name
+        return filename
+    return ""
+
+# Define the update logic for filename
+demo.input_components[0].change(update_filename, inputs=demo.input_components[0], outputs=demo.input_components[1])
+
 if __name__ == '__main__':
     demo.launch()