test lg

app.py

@@ -4,6 +4,7 @@ from PIL import Image
 import gradio as gr
 
 def save_image(image, path):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
     image.save(path)
 
 def load_image(path):
@@ -12,12 +13,11 @@ def load_image(path):
 def process_image(input_image):
     input_path = "./datasets/data/test/input_image.png"
     output_dir = "./results/demo/color_pix2pix/test_latest"
-    output_image_path = os.path.join(output_dir, "images", "fake_B.png")
+    output_image_path = os.path.join(output_dir, "images", "fake_B_rgb.png")
 
     # Save the input image
     save_image(input_image, input_path)
     
-
     cmd = [
         "python", "test.py",  # Command to run the test script
         "--dataroot", "./datasets/data",  # Adjust path as needed
@@ -33,42 +33,26 @@ def process_image(input_image):
         # Execute the command to process the image
         subprocess.run(cmd, check=True)
     except subprocess.CalledProcessError as e:
-        print(f"Error while running command: {e}")
-        return None
+        return f"Error while running command: {e}"
     
     # Check if the output directory exists
     if not os.path.exists(output_dir):
-        print(f"Error: Output directory {output_dir} does not exist.")
-        return None
+        return f"Error: Output directory {output_dir} does not exist."
 
     # After processing, load the output image from the results directory
     output_files = [f for f in os.listdir(os.path.join(output_dir, "images")) if f.endswith('fake_B_rgb.png')]
-    print(output_files)
     if not output_files:
-        print(f"Error: No output files found in {os.path.join(output_dir, 'images')}.")
-        return None
+        return f"Error: No output files found in {os.path.join(output_dir, 'images')}."
     
-    print(output_files)
-
-    newest_file = max(output_files, key=lambda f: os.path.getctime(os.path.join(output_dir, "images", f)))
-    output_image_path = os.path.join(output_dir, "images", newest_file)
-    if os.path.exists(output_image_path):
-        print(f"Output image saved at {output_image_path}")
-        return load_image(output_image_path)
-    else:
-        print(f"Error: Output image {output_image_path} not found.")
-        return None
+    return load_image(os.path.join(output_dir, "images", output_files[0]))
 
-# Define the Gradio interface
 iface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    live=True,
-    title="Pix2Pix Colorization",
-    description="Upload an image, which will be processed using Pix2Pix model and the output will be displayed."
+    title="Image Colorization",
+    description="Upload an image to colorize it using the model."
 )
 
-
-# Launch the app
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()