Update app.py

app.py

@@ -15,12 +15,12 @@ def string_to_binary_labels(s: str) -> list[int]:
     return bits
 
 def image_to_binary_labels(img: Image.Image, max_pixels: int = 256) -> list[int]:
-    img = img.convert("L")
+    img = img.convert("RGB")
     img.thumbnail((int(np.sqrt(max_pixels)), int(np.sqrt(max_pixels))))
     img_array = np.array(img)
     flat = img_array.flatten()
-    binarized = (flat > 128).astype(int)
-    return binarized.tolist()
+    binarized = [(pixel >> bit) & 1 for pixel in flat for bit in range(7, -1, -1)]
+    return binarized
 
 def binary_labels_to_image(binary_labels: list[int], width: int = None, height: int = None) -> Image.Image:
     total_pixels = len(binary_labels)
@@ -35,12 +35,27 @@ def binary_labels_to_image(binary_labels: list[int], width: int = None, height:
     img = Image.fromarray(image_array, mode='L')
     return img
 
-def colorize_grayscale(img: Image.Image) -> Image.Image:
-    img_array = np.array(img) / 255.0
-    colored_array = cm.viridis(img_array)[:, :, :3]
-    colored_array = (colored_array * 255).astype(np.uint8)
-    color_img = Image.fromarray(colored_array, mode='RGB')
-    return color_img
+def binary_labels_to_rgb_image(binary_labels: list[int], width: int = None, height: int = None) -> Image.Image:
+    total_pixels = len(binary_labels) // 24
+    if width is None or height is None:
+        side = int(np.ceil(np.sqrt(total_pixels)))
+        width = height = side
+    needed_pixels = width * height
+    needed_bits = needed_pixels * 24
+    if len(binary_labels) < needed_bits:
+        binary_labels += [0] * (needed_bits - len(binary_labels))
+    pixels = []
+    for i in range(0, needed_bits, 24):
+        r_bits = binary_labels[i:i+8]
+        g_bits = binary_labels[i+8:i+16]
+        b_bits = binary_labels[i+16:i+24]
+        r = sum(b << (7-j) for j, b in enumerate(r_bits))
+        g = sum(b << (7-j) for j, b in enumerate(g_bits))
+        b = sum(b << (7-j) for j, b in enumerate(b_bits))
+        pixels.append((r, g, b))
+    array = np.array(pixels, dtype=np.uint8).reshape((height, width, 3))
+    img = Image.fromarray(array, mode='RGB')
+    return img
 
 mutation_site_headers = [
     3244, 3297, 3350, 3399, 3455, 3509, 3562, 3614,
@@ -128,12 +143,13 @@ with tab2:
         )
 
         st.subheader("Reconstruct Image from Binary Labels")
-        if st.button("Reconstruct Image"):
-            reconstructed_img = binary_labels_to_image(binary_labels)
-            st.image(reconstructed_img, caption="Reconstructed Grayscale Image", use_column_width=True)
+        option = st.radio("Choose Reconstruction Mode", ["Grayscale", "True Color (RGB)"])
 
-            st.subheader("Colorized Image (Fake Color)")
-            colorized_img = colorize_grayscale(reconstructed_img)
-            st.image(colorized_img, caption="Colorized Reconstructed Image", use_column_width=True)
+        if st.button("Reconstruct Image"):
+            if option == "Grayscale":
+                reconstructed_img = binary_labels_to_image(binary_labels)
+            else:
+                reconstructed_img = binary_labels_to_rgb_image(binary_labels)
+            st.image(reconstructed_img, caption="Reconstructed Image", use_column_width=True)
 
-# Future: integrate DNA editor mapping for each mutation site here
+# Future: integrate DNA editor mapping for each mutation site here