Update app.py

app.py

@@ -300,9 +300,32 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
             
         processed_images = []
         for img in images:
-            img_np = img.cpu().permute(1, 2, 0).numpy()
+            # Convert to numpy and remove channel dimension (3,128,128 -> 128,128)
+            img_np = img.cpu().permute(1, 2, 0).mean(dim=-1).numpy()
+            
+            # Normalize to 0-255
             img_np = (img_np * 255).clip(0, 255).astype(np.uint8)
-            pil_img = Image.fromarray(img_np)
+            
+            # Apply bone colormap approximation (without matplotlib)
+            # This is a simplified version of the bone colormap
+            bone_cmap = np.array([
+                [0, 0, 0],       # Black
+                [25, 25, 51],     # Dark blue-gray
+                [50, 50, 102],    # Medium blue-gray
+                [76, 76, 153],    # Blue-gray
+                [101, 101, 204],  # Light blue-gray
+                [127, 127, 255],  # Very light blue
+                [153, 178, 255],  # Blue-white
+                [178, 204, 255],  # Light blue-white
+                [204, 229, 255],  # Very light blue-white
+                [229, 255, 255],  # Almost white
+                [255, 255, 255]   # White
+            ], dtype=np.uint8)
+            
+            # Apply colormap
+            colored = bone_cmap[(img_np * (len(bone_cmap)-1) / 255].astype(np.uint8)
+            pil_img = Image.fromarray(colored)
+            
             processed_images.append(pil_img)
         
         if num_images == 1: