Update app.py

app.py

@@ -3,10 +3,10 @@ import os
 import threading
 import subprocess
 import time
+import re
 from huggingface_hub import hf_hub_download
 from handwriting_api import InputData, validate_input
 from hand import Hand
-import cairosvg
 
 # Create img directory if it doesn't exist
 os.makedirs("img", exist_ok=True)
@@ -21,7 +21,8 @@ def generate_handwriting(
     bias=0.75, 
     color="#000000",
     stroke_width=2,
-    multiline=True
+    multiline=True,
+    transparent_background=True
 ):
     """Generate handwritten text using the model"""
     try:
@@ -74,30 +75,78 @@ def generate_handwriting(
         with open("img/output.svg", "r") as f:
             svg_content = f.read()
             
+        # If transparent background is requested, modify the SVG
+        if transparent_background:
+            # Remove the background rectangle or make it transparent
+            pattern = r'<rect[^>]*?fill="white"[^>]*?>'
+            if re.search(pattern, svg_content):
+                svg_content = re.sub(pattern, '', svg_content)
+            
+            # Write the modified SVG back
+            with open("img/output.svg", "w") as f:
+                f.write(svg_content)
+            
         return svg_content
     except Exception as e:
         return f"Error: {str(e)}"
 
 def export_to_png(svg_content):
-    """Convert SVG to transparent PNG using CairoSVG"""
+    """Convert SVG to transparent PNG using CairoSVG and Pillow for robust transparency"""
     try:
         import cairosvg
+        from PIL import Image
         
         if not svg_content or svg_content.startswith("Error:"):
             return None
-            
-        # Save SVG to a temporary file
+        
+        # Modify the SVG to ensure the background is transparent
+        # Remove any white background rectangle
+        pattern = r'<rect[^>]*?fill="white"[^>]*?>'
+        if re.search(pattern, svg_content):
+            svg_content = re.sub(pattern, '', svg_content)
+        
+        # Save the modified SVG to a temporary file
         with open("img/temp.svg", "w") as f:
             f.write(svg_content)
             
-        # Convert SVG to PNG with transparency
+        # Convert SVG to PNG with transparency using CairoSVG
         cairosvg.svg2png(
             url="img/temp.svg", 
-            write_to="img/output.png", 
+            write_to="img/output_temp.png", 
             scale=2.0,
             background_color="none"  # This ensures transparency
         )
         
+        # Additional processing with Pillow to ensure transparency
+        img = Image.open("img/output_temp.png")
+        
+        # Convert to RGBA if not already
+        if img.mode != 'RGBA':
+            img = img.convert('RGBA')
+        
+        # Create a transparent canvas
+        transparent_img = Image.new('RGBA', img.size, (0, 0, 0, 0))
+        
+        # Process the image data to ensure white is transparent
+        datas = img.getdata()
+        new_data = []
+        
+        for item in datas:
+            # If pixel is white or near-white, make it transparent
+            if item[0] > 240 and item[1] > 240 and item[2] > 240:
+                new_data.append((255, 255, 255, 0))  # Transparent
+            else:
+                new_data.append(item)  # Keep original color
+        
+        transparent_img.putdata(new_data)
+        transparent_img.save("img/output.png", "PNG")
+        
+        # Clean up the temporary file
+        try:
+            os.remove("img/output_temp.png")
+        except:
+            pass
+        
         return "img/output.png"
     except Exception as e:
         print(f"Error converting to PNG: {str(e)}")
@@ -239,12 +288,22 @@ if __name__ == "__main__":
     # Set port based on environment variable or default to 7860
     port = int(os.environ.get("PORT", 7860))
     
-    # Check if CairoSVG is installed
+    # Check if required packages are installed
+    missing_packages = []
     try:
         import cairosvg
-        print("CairoSVG is installed and ready for PNG export")
     except ImportError:
-        print("WARNING: CairoSVG is not installed. PNG export will not work.")
-        print("Please install it using: pip install cairosvg")
+        missing_packages.append("cairosvg")
+    
+    try:
+        from PIL import Image
+    except ImportError:
+        missing_packages.append("pillow")
+    
+    if missing_packages:
+        print(f"WARNING: The following packages are missing and required for transparent PNG export: {', '.join(missing_packages)}")
+        print("Please install them using: pip install " + " ".join(missing_packages))
+    else:
+        print("All required packages are installed and ready for transparent PNG export")
     
     demo.launch(server_name="0.0.0.0", server_port=port)