Update app.py

app.py

@@ -1,21 +1,22 @@
 import os
-
 import torch
-import torch.nn.functional as F
+import cv2
 from torchvision.transforms import Compose, ToTensor, Resize, Normalize, ConvertImageDtype
-
+from PIL import Image
 import numpy as np
-import cv2
 
 import gradio as gr
-from huggingface_hub import hf_hub_download
-
 from model import IAT
 
-
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
 
+def tensor_to_numpy(tensor):
+    tensor = tensor.detach().cpu().numpy()
+    if tensor.ndim == 3 and tensor.shape[0] == 3:  # Convert CHW to HWC
+        tensor = tensor.transpose(1, 2, 0)
+    tensor = np.clip(tensor * 255, 0, 255).astype(np.uint8)  # Ensure the output is uint8
+    return tensor
 
 def dark_inference(img):
     model = IAT()
@@ -34,9 +35,11 @@ def dark_inference(img):
     input_img = transform(img)
     print(f'Image shape: {input_img.shape}')
 
-    enhanced_img = model(input_img.unsqueeze(0))
-    return enhanced_img[0].permute(1, 2, 0).detach().numpy()
-
+    with torch.no_grad():
+        enhanced_img = model(input_img.unsqueeze(0))
+    
+    result_img = tensor_to_numpy(enhanced_img[0])
+    return result_img
 
 def exposure_inference(img):
     model = IAT()
@@ -49,14 +52,17 @@ def exposure_inference(img):
     transform = Compose([
         ToTensor(), 
         Resize(384),
+        Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 
         ConvertImageDtype(torch.float) 
     ])
     input_img = transform(img)
     print(f'Image shape: {input_img.shape}')
 
-    enhanced_img = model(input_img.unsqueeze(0))
-    return enhanced_img[0].permute(1, 2, 0).detach().numpy()
-
+    with torch.no_grad():
+        enhanced_img = model(input_img.unsqueeze(0))
+    
+    result_img = tensor_to_numpy(enhanced_img[0])
+    return result_img
 
 demo = gr.Blocks()
 with demo:
@@ -77,7 +83,7 @@ with demo:
                     with gr.Row():
                         exposure_button = gr.Button('Exposure Correction')
                 with gr.Column():
-                    res_image = gr.Image(type='numpy', label='Resutls')
+                    res_image = gr.Image(type='numpy', label='Results')
         with gr.Row():
             dark_example_images = gr.Dataset(
                 components=[input_image], 
@@ -100,4 +106,4 @@ with demo:
     dark_example_images.click(fn=set_example_image, inputs=dark_example_images, outputs=dark_example_images.components)
     exposure_example_images.click(fn=set_example_image, inputs=exposure_example_images, outputs=exposure_example_images.components)
 
-demo.launch(enable_queue=True)
+demo.launch(enable_queue=True)