Update app.py

app.py

@@ -8,6 +8,8 @@ import os
 from threading import Event
 import traceback
 import cv2  # Added for bilateral filtering
+import matplotlib.pyplot as plt
+from io import BytesIO
 
 # Constants
 IMG_SIZE = 128
@@ -300,13 +302,31 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
             
         processed_images = []
         for img in images:
-            # Convert to grayscale (X-rays are naturally grayscale)
-            # Take the mean across RGB channels and convert to uint8
+            # Convert to grayscale and apply bone colormap
             img_np = img.cpu().permute(1, 2, 0).mean(dim=-1).numpy()
-            img_np = (img_np * 255).clip(0, 255).astype(np.uint8)
             
-            # Create PIL image in grayscale mode
-            pil_img = Image.fromarray(img_np, mode='L')
+            # Normalize to 0-1
+            img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)
+            
+            # Apply additional sharpening with OpenCV
+            img_np_uint8 = (img_np * 255).astype(np.uint8)
+            
+            # Apply unsharp mask for additional sharpness
+            blurred = cv2.GaussianBlur(img_np_uint8, (0, 0), 2.0)
+            sharpened = cv2.addWeighted(img_np_uint8, 1.5, blurred, -0.5, 0)
+            
+            # Apply bone colormap using matplotlib
+            plt.figure(figsize=(2.56, 2.56), dpi=50)
+            plt.imshow(sharpened, cmap='bone')
+            plt.axis('off')
+            plt.tight_layout(pad=0)
+            
+            # Save to buffer
+            buf = BytesIO()
+            plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=50)
+            buf.seek(0)
+            pil_img = Image.open(buf)
+            plt.close()
             
             processed_images.append(pil_img)