Update app.py

app.py

@@ -11,15 +11,7 @@ device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cp
 pil_to_tensor = transforms.ToTensor()
 
 # 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'
-    }
+image_to_weights = {'./weights/flolv2_UHDLL.pt','./weights/flolv2_all_111439.pt'}
 
 # Initial model setup (without weights)
 model = create_model()
@@ -32,12 +24,12 @@ def load_img(filename):
 def process_img(image, filename):
     # Select the correct weight file based on the image filename
 
-    filename = image.name.split("/")[-1]
-    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)
+    # filename = image.name.split("/")[-1]
+    # if filename in image_to_weights:
+    model_path = image_to_weights[image_to_weights[1]]
+    checkpoints = torch.load(model_path, map_location=device)
+    model.load_state_dict(checkpoints['params'])
+    model.to(device)
 
     img = np.array(image)
     img = img / 255.  # Normalize to [0, 1]
@@ -64,12 +56,11 @@ Due to the GPU memory limitations, the app might crash if you feed a high-resolu
 <br>
 '''
 
-examples = [['images/425_UHD_LL.JPG', '425_UHD_LL.JPG'],
-            ['images/low00772.png', 'low00772.png'], 
-            ['images/low00723.png', 'low00723.png'], 
-            ['images/low00748.png', 'low00748.png'], 
-            ['images/1778_UHD_LL.JPG', '1778_UHD_LL.JPG'],
-            ['images/1791_UHD_LL.JPG', '1791_UHD_LL.JPG']]
+examples = [
+            ['images/low00772.png'], 
+            ['images/low00723.png'], 
+            ['images/low00748.png'], 
+           ]
 
 css = """
     .image-frame img, .image-container img {
@@ -81,22 +72,17 @@ css = """
 
 
 
-with gr.Blocks() as demo:
-    with gr.Row():
-        input_image = gr.Image(type='pil', label='input', interactive = True)
-        filename_output = gr.Textbox(label="Image Filename", interactive=False)
-
-    # Define the output
-    output_image = gr.Image(type='pil', label='output')
-
-    # Define the interaction flow: when the image is uploaded, update the filename
-    input_image.change(update_filename, inputs=input_image, outputs=filename_output)
-    
-    # Set the function to process the image with the filename
-    input_image.change(process_img, inputs=[input_image, filename_output], outputs=output_image)
-
-    # Provide examples
-    demo.examples = examples
+demo = gr.Interface(
+    fn = process_img,
+    inputs = [
+            gr.Image(type = 'pil', label = 'input')
+    ],
+    outputs = [gr.Image(type='pil', label = 'output')],
+    title = title,
+    description = description,
+    examples = examples,
+    css = css
+)
 
 if __name__ == '__main__':
     demo.launch()