Update app.py

app.py

@@ -7,6 +7,9 @@ import math
 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
@@ -153,8 +156,10 @@ class DiffusionModel(nn.Module):
 
     @torch.no_grad()
     def sample(self, num_images, img_size, num_classes, labels, device, progress_callback=None):
+        # Start with random noise
         x_t = torch.randn(num_images, 3, img_size, img_size).to(device)
 
+        # Label handling (one-hot if needed)
         if labels.ndim == 1:
             labels_one_hot = torch.zeros(num_images, num_classes).to(device)
             labels_one_hot[torch.arange(num_images), labels] = 1
@@ -162,6 +167,7 @@ class DiffusionModel(nn.Module):
         else:
             labels = labels.to(device)
 
+        # REVERTED SAMPLING LOOP WITH NOISE REDUCTION
         for t in reversed(range(self.timesteps)):
             if cancel_event.is_set():
                 return None
@@ -169,6 +175,7 @@ class DiffusionModel(nn.Module):
             t_tensor = torch.full((num_images,), t, device=device, dtype=torch.float)
             predicted_noise = self.model(x_t, labels, t_tensor)
 
+            # Calculate coefficients
             beta_t = self.betas[t].to(device)
             alpha_t = self.alphas[t].to(device)
             alpha_bar_t = self.alpha_bars[t].to(device)
@@ -176,8 +183,9 @@ class DiffusionModel(nn.Module):
             mean = (1 / torch.sqrt(alpha_t)) * (x_t - (beta_t / torch.sqrt(1 - alpha_bar_t)) * predicted_noise)
             variance = beta_t
 
+            # Reduced noise injection with lower multiplier
             if t > 0:
-                noise = torch.randn_like(x_t)
+                noise = torch.randn_like(x_t) * 0.8  # Reduced noise by 20%
             else:
                 noise = torch.zeros_like(x_t)
 
@@ -186,14 +194,34 @@ class DiffusionModel(nn.Module):
             if progress_callback:
                 progress_callback((self.timesteps - t) / self.timesteps)
 
+        # Clamp and denormalize
         x_0 = torch.clamp(x_t, -1., 1.)
-
-        # Normalization
         mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
         std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
         x_0 = std * x_0 + mean
         x_0 = torch.clamp(x_0, 0., 1.)
 
+        # ENHANCED SHARPENING
+        # First apply mild bilateral filtering to reduce noise while preserving edges
+        x_np = x_0.cpu().permute(0, 2, 3, 1).numpy()
+        filtered = []
+        for img in x_np:
+            img = (img * 255).astype(np.uint8)
+            filtered_img = cv2.bilateralFilter(img, d=5, sigmaColor=15, sigmaSpace=15)
+            filtered.append(filtered_img / 255.0)
+        x_0 = torch.tensor(np.array(filtered), device=device).permute(0, 3, 1, 2)
+        
+        # Then apply stronger unsharp masking
+        kernel = torch.ones(3, 1, 5, 5, device=device) / 75
+        kernel = kernel.to(x_0.dtype)
+        blurred = torch.nn.functional.conv2d(
+            x_0,
+            kernel,
+            padding=2,
+            groups=3
+        )
+        x_0 = torch.clamp(1.5 * x_0 - 0.5 * blurred, 0., 1.)  # Increased sharpening factor
+
         return x_0
 
 def load_model(model_path, device):
@@ -274,9 +302,34 @@ 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()
-            img_np = (img_np * 255).clip(0, 255).astype(np.uint8)
-            pil_img = Image.fromarray(img_np)
+            # Convert to grayscale and apply bone colormap
+            img_np = img.cpu().permute(1, 2, 0).mean(dim=-1).numpy()
+            
+            # 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 - FIXED APPROACH
+            # Create a simple bone-like colormap manually to avoid matplotlib issues
+            sharpened_normalized = sharpened / 255.0
+            # Simulate bone colormap: black to white with blueish tones
+            r = np.clip(sharpened_normalized * 1.2 - 0.1, 0, 1)
+            g = np.clip(sharpened_normalized * 1.1 - 0.05, 0, 1)
+            b = np.clip(sharpened_normalized * 1.0 + 0.1, 0, 1)
+            
+            # Combine channels and convert to uint8
+            bone_colored = np.stack([r, g, b], axis=-1)
+            bone_colored_uint8 = (bone_colored * 255).astype(np.uint8)
+            
+            # Create PIL image
+            pil_img = Image.fromarray(bone_colored_uint8)
+            
             processed_images.append(pil_img)
         
         if num_images == 1: