Update app.py

app.py

@@ -10,6 +10,12 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 weights_path = "unet_model.pth"
 model = load_model(weights_path, device)
 
+def hex_to_rgb(hex_color):
+    hex_color = hex_color.lstrip("#")  
+    if len(hex_color) == 3:  
+        hex_color = "".join([c * 2 for c in hex_color])
+    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
+
 def process_image(image, text, font_size, text_color):
     image = image.convert("RGB")
     print(f"image: {image}")
@@ -20,7 +26,7 @@ def process_image(image, text, font_size, text_color):
     font = ImageFont.truetype(font_path, font_size)
     text_position = (50, 50)
     # text_color = (0, 0, 0)
-    text_color = tuple(int(text_color[i:i+2], 16) for i in (1, 3, 5))
+    text_color = hex_to_rgb(text_color)
     draw.text(text_position, text, fill=text_color, font=font)
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")