Update app.py

app.py

@@ -33,7 +33,15 @@ def predict(image):
 
         # پس‌پردازش خروجی
         output_image = output.squeeze().cpu().numpy()
-        output_image = output_image.transpose(1, 2, 0)  # تبدیل از CHW به HWC
+
+        # اگر خروجی تک کاناله است، به RGB تبدیل می‌شود
+        if output_image.ndim == 2:  # فقط در صورتی که تک کاناله است
+            output_image = np.stack([output_image] * 3, axis=-1)
+        
+        elif output_image.shape[0] == 1:  # اگر کانال اول 1 است، آن را به RGB تبدیل کنید
+            output_image = np.tile(output_image, (3, 1, 1))
+            output_image = output_image.transpose(1, 2, 0)
+
         output_image = (output_image * 255).astype(np.uint8)
         output_image = Image.fromarray(output_image)