data-science-mango commited on
Commit
b3460e8
·
verified ·
1 Parent(s): 5db557b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -0
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ from PIL import Image
4
+ import tempfile
5
+ import os
6
+
7
+ def generate_binary_mask(image_data):
8
+ if image_data is None or "layers" not in image_data or not image_data["layers"]:
9
+ raise gr.Error("Please draw a mask before generating!")
10
+
11
+ mask = image_data["layers"][0]
12
+ mask_array = np.array(mask)
13
+
14
+ if np.all(mask_array < 10):
15
+ raise gr.Error("The mask is empty! Please draw something.")
16
+
17
+ # Binary mask logic
18
+ is_black = np.all(mask_array < 10, axis=2)
19
+ binary_mask = Image.fromarray(((~is_black) * 255).astype(np.uint8))
20
+
21
+ # Save to temporary file
22
+ temp_dir = tempfile.mkdtemp()
23
+ output_path = os.path.join(temp_dir, "binary_mask.png")
24
+ binary_mask.save(output_path)
25
+
26
+ return binary_mask, output_path
27
+
28
+ with gr.Blocks() as app:
29
+ with gr.Row():
30
+ with gr.Column():
31
+ image_input = gr.ImageMask(
32
+ label="Upload or Paste Image, then draw mask",
33
+ type="pil",
34
+ height=None,
35
+ width=None
36
+ )
37
+ generate_btn = gr.Button("Generate Mask")
38
+
39
+ with gr.Column():
40
+ mask_preview = gr.Image(label="Mask Preview", type="pil")
41
+ output_file = gr.File(label="Download Mask")
42
+
43
+ generate_btn.click(fn=generate_binary_mask, inputs=image_input, outputs=[mask_preview, output_file])
44
+
45
+ app.launch()