Boni98 commited on
Commit
4179885
·
1 Parent(s): 5e86714
Files changed (2) hide show
  1. app.py +19 -123
  2. old_app.py +135 -0
app.py CHANGED
@@ -1,135 +1,31 @@
1
  import gradio as gr
2
- import numpy as np
3
- from PIL import Image
4
- import tempfile
5
- import os
6
 
7
- # Function to create the binary mask from the ImageEditor's output
8
- def create_binary_mask(im_dict):
9
- """
10
- Generates a binary mask from the drawing layer of the gr.ImageEditor output.
11
 
12
- Args:
13
- im_dict (dict): The dictionary output from gr.ImageEditor, containing
14
- 'background', 'layers', and 'composite'.
15
 
16
- Returns:
17
- tuple: A tuple containing:
18
- - np.ndarray: The binary mask image (H, W) as a NumPy array (0 or 255).
19
- - str or None: The filepath to the saved PNG mask for download, or None if no mask generated.
20
- """
21
- if im_dict is None or im_dict["background"] is None:
22
- print("No background image found.")
23
- # Return a small blank placeholder and None for the file path
24
- blank_preview = np.zeros((768, 1024), dtype=np.uint8)
25
- return blank_preview, None
26
 
27
- background_img = im_dict["background"]
28
- h, w, _ = background_img.shape # Get original dimensions (Height, Width, Channels)
29
- print(f"Original image dimensions: H={h}, W={w}")
30
 
31
- # Check if any drawing layer exists and is not None
32
- if not im_dict["layers"] or im_dict["layers"][0] is None:
33
- print("No drawing layer found. Generating blank mask.")
34
- # Nothing drawn yet, return a black mask of the original size
35
- mask = np.zeros((h, w), dtype=np.uint8)
36
- filepath = None # No file to download as nothing was drawn
37
- else:
38
- # Use the first layer (index 0) which usually contains the drawing
39
- layer = im_dict["layers"][0]
40
- print(f"Drawing layer dimensions: H={layer.shape[0]}, W={layer.shape[1]}")
41
 
42
- # Ensure layer dimensions match background (Gradio ImageEditor usually handles this)
43
- if layer.shape[0] != h or layer.shape[1] != w:
44
- print(f"Warning: Layer size ({layer.shape[0]}x{layer.shape[1]}) doesn't match background ({h}x{w}). This shouldn't happen.")
45
- # Handle potential mismatch if necessary, though unlikely with default editor behavior
46
- # For now, proceed assuming they match or the layer is the correct reference
47
-
48
- # Layer is RGBA, extract the Alpha channel (index 3)
49
- alpha_channel = layer[:, :, 3]
50
-
51
- # Create binary mask: white (255) where alpha > 0 (drawn), black (0) otherwise
52
- mask = np.where(alpha_channel > 0, 255, 0).astype(np.uint8)
53
- print(f"Generated binary mask dimensions: H={mask.shape[0]}, W={mask.shape[1]}")
54
-
55
- # Save the mask to a temporary PNG file for download
56
- try:
57
- # Create a temporary file path
58
- with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
59
- filepath = tmpfile.name
60
-
61
- # Save the NumPy array as a PNG image using PIL
62
- pil_image = Image.fromarray(mask)
63
- pil_image.save(filepath, format="PNG")
64
- print(f"Mask saved temporarily to: {filepath}")
65
-
66
- except Exception as e:
67
- print(f"Error saving mask to temporary file: {e}")
68
- filepath = None # Indicate failure to save
69
- # Return a blank mask in case of saving error
70
- mask = np.zeros((h, w), dtype=np.uint8)
71
-
72
- # Return the mask NumPy array for preview and the filepath for download
73
- # The DownloadButton component will become active/functional if filepath is not None
74
- return mask, filepath
75
-
76
- # --- Gradio App Layout ---
77
  with gr.Blocks() as demo:
78
- gr.Markdown("## Binary Mask Generator")
79
- gr.Markdown(
80
- "Upload or paste an image. Use the brush tool (select it!) to draw the area "
81
- "you want to mask. Click 'Generate Mask' to see the result and download it."
82
- )
83
-
84
  with gr.Row():
85
- # --- Left Column ---
86
- with gr.Column(scale=1): # Adjust scale as needed
87
- image_editor = gr.ImageEditor(
88
- label="Draw on Image",
89
- # type="numpy" is essential for processing layers
90
- type="numpy",
91
- # DON'T set crop_size, height, or width to keep original dimensions
92
- # sources allow upload, paste, webcam etc.
93
- sources=["upload"],
94
- # Set a default brush for clarity (optional, but helpful)
95
- brush=gr.Brush(colors=["#FF0000"], color_mode="fixed"), # Red fixed brush
96
- interactive=True,
97
- canvas_size=(768, 1024)
98
- )
99
- generate_button = gr.Button("Generate Mask", variant="primary")
100
-
101
- # --- Right Column ---
102
- with gr.Column(scale=1): # Adjust scale as needed
103
- mask_preview = gr.Image(
104
- label="Binary Mask Preview",
105
- # Use numpy for consistency, PIL would also work
106
- type="numpy",
107
- interactive=False, # Preview is not interactive
108
- )
109
- # Download button - its value (the file path) is set by the function's output
110
- download_button = gr.DownloadButton(
111
- label="Download Mask (PNG)",
112
- interactive=True, # Button starts interactive
113
- )
114
 
115
- # --- Event Handling ---
116
- generate_button.click(
117
- fn=create_binary_mask,
118
- inputs=[image_editor],
119
- # Output 1 goes to mask_preview (image data)
120
- # Output 2 goes to download_button (file path for the 'value' argument)
121
- outputs=[mask_preview, download_button]
122
- )
123
-
124
- # --- Launch the App ---
125
  if __name__ == "__main__":
126
- # Cleaning up old temp files on startup (optional but good practice)
127
- temp_dir = tempfile.gettempdir()
128
- for item in os.listdir(temp_dir):
129
- if item.endswith(".png") and item.startswith("tmp"): # Be specific to avoid deleting wrong files
130
- try:
131
- os.remove(os.path.join(temp_dir, item))
132
- except Exception:
133
- pass # Ignore if file is locked etc.
134
-
135
  demo.launch()
 
1
  import gradio as gr
2
+ import time
 
 
 
3
 
 
 
 
 
4
 
5
+ def sleep(im):
6
+ time.sleep(5)
7
+ return [im["background"], im["layers"][0], im["layers"][1], im["composite"]]
8
 
 
 
 
 
 
 
 
 
 
 
9
 
10
+ def predict(im):
11
+ return im["composite"]
 
12
 
 
 
 
 
 
 
 
 
 
 
13
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  with gr.Blocks() as demo:
 
 
 
 
 
 
15
  with gr.Row():
16
+ im = gr.ImageEditor(
17
+ type="numpy",
18
+ crop_size="1:1",
19
+ )
20
+ im_preview = gr.Image()
21
+ n_upload = gr.Number(0, label="Number of upload events", step=1)
22
+ n_change = gr.Number(0, label="Number of change events", step=1)
23
+ n_input = gr.Number(0, label="Number of input events", step=1)
24
+
25
+ im.upload(lambda x: x + 1, outputs=n_upload, inputs=n_upload)
26
+ im.change(lambda x: x + 1, outputs=n_change, inputs=n_change)
27
+ im.input(lambda x: x + 1, outputs=n_input, inputs=n_input)
28
+ im.change(predict, outputs=im_preview, inputs=im, show_progress="hidden")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
 
 
 
 
 
 
 
 
 
 
 
30
  if __name__ == "__main__":
 
 
 
 
 
 
 
 
 
31
  demo.launch()
old_app.py ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ from PIL import Image
4
+ import tempfile
5
+ import os
6
+
7
+ # Function to create the binary mask from the ImageEditor's output
8
+ def create_binary_mask(im_dict):
9
+ """
10
+ Generates a binary mask from the drawing layer of the gr.ImageEditor output.
11
+
12
+ Args:
13
+ im_dict (dict): The dictionary output from gr.ImageEditor, containing
14
+ 'background', 'layers', and 'composite'.
15
+
16
+ Returns:
17
+ tuple: A tuple containing:
18
+ - np.ndarray: The binary mask image (H, W) as a NumPy array (0 or 255).
19
+ - str or None: The filepath to the saved PNG mask for download, or None if no mask generated.
20
+ """
21
+ if im_dict is None or im_dict["background"] is None:
22
+ print("No background image found.")
23
+ # Return a small blank placeholder and None for the file path
24
+ blank_preview = np.zeros((768, 1024), dtype=np.uint8)
25
+ return blank_preview, None
26
+
27
+ background_img = im_dict["background"]
28
+ h, w, _ = background_img.shape # Get original dimensions (Height, Width, Channels)
29
+ print(f"Original image dimensions: H={h}, W={w}")
30
+
31
+ # Check if any drawing layer exists and is not None
32
+ if not im_dict["layers"] or im_dict["layers"][0] is None:
33
+ print("No drawing layer found. Generating blank mask.")
34
+ # Nothing drawn yet, return a black mask of the original size
35
+ mask = np.zeros((h, w), dtype=np.uint8)
36
+ filepath = None # No file to download as nothing was drawn
37
+ else:
38
+ # Use the first layer (index 0) which usually contains the drawing
39
+ layer = im_dict["layers"][0]
40
+ print(f"Drawing layer dimensions: H={layer.shape[0]}, W={layer.shape[1]}")
41
+
42
+ # Ensure layer dimensions match background (Gradio ImageEditor usually handles this)
43
+ if layer.shape[0] != h or layer.shape[1] != w:
44
+ print(f"Warning: Layer size ({layer.shape[0]}x{layer.shape[1]}) doesn't match background ({h}x{w}). This shouldn't happen.")
45
+ # Handle potential mismatch if necessary, though unlikely with default editor behavior
46
+ # For now, proceed assuming they match or the layer is the correct reference
47
+
48
+ # Layer is RGBA, extract the Alpha channel (index 3)
49
+ alpha_channel = layer[:, :, 3]
50
+
51
+ # Create binary mask: white (255) where alpha > 0 (drawn), black (0) otherwise
52
+ mask = np.where(alpha_channel > 0, 255, 0).astype(np.uint8)
53
+ print(f"Generated binary mask dimensions: H={mask.shape[0]}, W={mask.shape[1]}")
54
+
55
+ # Save the mask to a temporary PNG file for download
56
+ try:
57
+ # Create a temporary file path
58
+ with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
59
+ filepath = tmpfile.name
60
+
61
+ # Save the NumPy array as a PNG image using PIL
62
+ pil_image = Image.fromarray(mask)
63
+ pil_image.save(filepath, format="PNG")
64
+ print(f"Mask saved temporarily to: {filepath}")
65
+
66
+ except Exception as e:
67
+ print(f"Error saving mask to temporary file: {e}")
68
+ filepath = None # Indicate failure to save
69
+ # Return a blank mask in case of saving error
70
+ mask = np.zeros((h, w), dtype=np.uint8)
71
+
72
+ # Return the mask NumPy array for preview and the filepath for download
73
+ # The DownloadButton component will become active/functional if filepath is not None
74
+ return mask, filepath
75
+
76
+ # --- Gradio App Layout ---
77
+ with gr.Blocks() as demo:
78
+ gr.Markdown("## Binary Mask Generator")
79
+ gr.Markdown(
80
+ "Upload or paste an image. Use the brush tool (select it!) to draw the area "
81
+ "you want to mask. Click 'Generate Mask' to see the result and download it."
82
+ )
83
+
84
+ with gr.Row():
85
+ # --- Left Column ---
86
+ with gr.Column(scale=1): # Adjust scale as needed
87
+ image_editor = gr.ImageEditor(
88
+ label="Draw on Image",
89
+ # type="numpy" is essential for processing layers
90
+ type="numpy",
91
+ # DON'T set crop_size, height, or width to keep original dimensions
92
+ # sources allow upload, paste, webcam etc.
93
+ sources=["upload"],
94
+ # Set a default brush for clarity (optional, but helpful)
95
+ brush=gr.Brush(colors=["#FF0000"], color_mode="fixed"), # Red fixed brush
96
+ interactive=True,
97
+ canvas_size=(768, 1024)
98
+ )
99
+ generate_button = gr.Button("Generate Mask", variant="primary")
100
+
101
+ # --- Right Column ---
102
+ with gr.Column(scale=1): # Adjust scale as needed
103
+ mask_preview = gr.Image(
104
+ label="Binary Mask Preview",
105
+ # Use numpy for consistency, PIL would also work
106
+ type="numpy",
107
+ interactive=False, # Preview is not interactive
108
+ )
109
+ # Download button - its value (the file path) is set by the function's output
110
+ download_button = gr.DownloadButton(
111
+ label="Download Mask (PNG)",
112
+ interactive=True, # Button starts interactive
113
+ )
114
+
115
+ # --- Event Handling ---
116
+ generate_button.click(
117
+ fn=create_binary_mask,
118
+ inputs=[image_editor],
119
+ # Output 1 goes to mask_preview (image data)
120
+ # Output 2 goes to download_button (file path for the 'value' argument)
121
+ outputs=[mask_preview, download_button]
122
+ )
123
+
124
+ # --- Launch the App ---
125
+ if __name__ == "__main__":
126
+ # Cleaning up old temp files on startup (optional but good practice)
127
+ temp_dir = tempfile.gettempdir()
128
+ for item in os.listdir(temp_dir):
129
+ if item.endswith(".png") and item.startswith("tmp"): # Be specific to avoid deleting wrong files
130
+ try:
131
+ os.remove(os.path.join(temp_dir, item))
132
+ except Exception:
133
+ pass # Ignore if file is locked etc.
134
+
135
+ demo.launch()