errors fixed + preview updated

app.py

@@ -1,23 +1,34 @@
-import keras
-import numpy as np
 import gradio as gr
-from tibblingai import wta
 from PIL import Image, ImageOps
-import matplotlib.pyplot as plt 
-
-
+import numpy as np
+import matplotlib.pyplot as plt  # needed for plt.imread
 def load_sample_images():
     sample_paths = ["sample1.jpg", "sample2.jpg"]  # Must be in the same folder as your script
     original_imgs = [Image.open(path) for path in sample_paths]
     processed_imgs = [process_image(path) for path in sample_paths]
     return gr.update(visible=True, value=original_imgs), gr.update(visible=True, value=processed_imgs)
 
+
+
 # Image processing function (you can replace this)
 def process_image(img: Image.Image) -> Image.Image:
-    img_tensor = wta.wta(plt.imread(img)).numpy()
-    print(img_tensor.shape)
-    # image_array = np.squeeze(img_tensor)
-    return img_tensor[0]
+    # Read image using matplotlib
+    img_np = plt.imread(img)
+    if img_np.ndim == 3 and img_np.shape[2] > 3:
+        img_np = img_np[:, :, :3]
+
+    # Convert to grayscale for original image display
+    gray_img = ImageOps.grayscale(Image.fromarray((img_np * 255).astype(np.uint8)))
+
+    # Run your WTA processing (dummy if not available)
+    # Replace this line with actual WTA processing
+    img_tensor = wta.wta(img_np).numpy()  # Assuming returns shape (1, H, W)
+
+    # Convert processed image to inferno colormap
+    inferno_colored = plt.cm.inferno(img_tensor[0])
+    inferno_img = Image.fromarray((inferno_colored[:, :, :3] * 255).astype(np.uint8))
+
+    return gray_img, inferno_img
 
 # Function to process uploaded images
 def process_images(images):
@@ -28,8 +39,9 @@ def process_images(images):
     processed_imgs = []
 
     for img in images:
-        original_imgs.append(img)
-        processed_imgs.append(process_image(img))
+        original_gray_scale_img, inferno_img = process_image(img)
+        original_imgs.append(original_gray_scale_img)
+        processed_imgs.append(inferno_img)
 
     return gr.update(visible=True, value=original_imgs), gr.update(visible=True, value=processed_imgs)
 
@@ -60,4 +72,4 @@ with gr.Blocks() as demo:
 
     # example_files = [["sample1.jpg", "sample2.jpg"], ["sample2.jpg"]]
     # gr.Examples(examples=example_files, inputs=[file_input], label="Try one of our example samples")
-demo.launch()
+demo.launch(debug=True)