Update app.py

app.py

@@ -6,12 +6,13 @@ from PIL import Image
 import numpy as np
 
 import gradio as gr
-from model import IAT
+from iatenhancement.model import IAT  # Ensure the correct import path
 
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
 
 def tensor_to_numpy(tensor):
+    print("Converting tensor to numpy array...")
     tensor = tensor.detach().cpu().numpy()
     if tensor.ndim == 3 and tensor.shape[0] == 3:  # Convert CHW to HWC
         tensor = tensor.transpose(1, 2, 0)
@@ -19,6 +20,7 @@ def tensor_to_numpy(tensor):
     return tensor
 
 def dark_inference(img):
+    print("Starting dark inference...")
     model = IAT()
     checkpoint_file_path = './checkpoint/best_Epoch_lol.pth'
     state_dict = torch.load(checkpoint_file_path, map_location='cpu')
@@ -33,15 +35,17 @@ def dark_inference(img):
         ConvertImageDtype(torch.float) 
     ])
     input_img = transform(img)
-    print(f'Image shape: {input_img.shape}')
+    print(f'Image shape after transform: {input_img.shape}')
 
     with torch.no_grad():
         enhanced_img = model(input_img.unsqueeze(0))
     
     result_img = tensor_to_numpy(enhanced_img[0])
+    print("Dark inference completed.")
     return result_img
 
 def exposure_inference(img):
+    print("Starting exposure inference...")
     model = IAT()
     checkpoint_file_path = './checkpoint/best_Epoch_exposure.pth'
     state_dict = torch.load(checkpoint_file_path, map_location='cpu')
@@ -56,12 +60,13 @@ def exposure_inference(img):
         ConvertImageDtype(torch.float) 
     ])
     input_img = transform(img)
-    print(f'Image shape: {input_img.shape}')
+    print(f'Image shape after transform: {input_img.shape}')
 
     with torch.no_grad():
         enhanced_img = model(input_img.unsqueeze(0))
     
     result_img = tensor_to_numpy(enhanced_img[0])
+    print("Exposure inference completed.")
     return result_img
 
 demo = gr.Blocks()