awacke1 commited on
Commit
7cd8e6e
·
verified ·
1 Parent(s): 3f97e41

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +107 -40
app.py CHANGED
@@ -5,12 +5,17 @@ import gradio as gr
5
  from PIL import Image, ImageFilter, ImageOps, ImageChops
6
  import torchvision.transforms as transforms
7
  import os
 
8
  import pathlib
9
 
10
  # --- ⚙️ Configuration ---
 
11
  output_dir = "outputs"
12
  os.makedirs(output_dir, exist_ok=True)
13
 
 
 
 
14
  # --- 🎨 Filters ---
15
  FILTERS = {
16
  "Standard": "📄", "Invert": "⚫⚪", "Blur": "🌫️", "Sharpen": "🔪", "Contour": "🗺️",
@@ -105,11 +110,13 @@ def apply_filter(line_img, filter_name, original_img):
105
  return Image.fromarray(noisy_array)
106
  return line_img
107
 
108
- # --- 🖼️ Main Processing Function ---
109
- def predict(input_img_path, line_style, filter_choice):
110
  if not model1 or not model2:
111
  raise gr.Error("Models are not loaded. Please check for 'model.pth' and 'model2.pth'.")
112
-
 
 
113
  filter_name = filter_choice.split(" ", 1)[1]
114
  original_img = Image.open(input_img_path).convert('RGB')
115
 
@@ -128,50 +135,110 @@ def predict(input_img_path, line_style, filter_choice):
128
 
129
  final_image = apply_filter(line_drawing_full_res, filter_name, original_img)
130
 
131
- # --- 💾 Save the output image ---
132
  base_name = pathlib.Path(input_img_path).stem
133
  output_filename = f"{base_name}_{filter_name}.png"
134
  output_filepath = os.path.join(output_dir, output_filename)
135
  final_image.save(output_filepath)
136
 
137
- return final_image
 
 
 
 
138
 
139
  # --- 🚀 Gradio UI Setup ---
140
  title = "🖌️ Image to Line Art with Creative Filters"
141
- description = "Upload an image, choose a line style (Complex or Simple), and select a filter from the dropdown to transform your picture. Results are saved in the 'outputs' folder."
142
-
143
- filter_choices = [f"{emoji} {name}" for name, emoji in FILTERS.items()]
144
-
145
- # --- New Curated Examples Section ---
146
- examples = []
147
- example_images = [f"{i:02d}.jpeg" for i in range(1, 11)]
148
- # A selection of 6 interesting filters to demonstrate
149
- demo_filters = ["🗺️ Contour", "🔵⚫ BlueOnBlack", "✖️ Multiply", "🏞️ Emboss", "🔪 Sharpen", "❄️ Noise"]
150
-
151
- # Create one example for each of the 10 image files, cycling through the demo filters
152
- for i, img_file in enumerate(example_images):
153
- if os.path.exists(img_file):
154
- # Use modulo to cycle through the 6 demo filters for the 10 images
155
- chosen_filter = demo_filters[i % len(demo_filters)]
156
- examples.append([img_file, 'Simple Lines', chosen_filter])
157
-
158
- if not examples:
159
- print("⚠️ Warning: No example images ('01.jpeg' to '10.jpeg') found. Examples will be empty.")
160
-
161
- # Reverted to the simpler and more stable gr.Interface
162
- iface = gr.Interface(
163
- fn=predict,
164
- inputs=[
165
- gr.Image(type='filepath', label="Upload Image"),
166
- gr.Radio(['Complex Lines', 'Simple Lines'], label='Line Style', value='Simple Lines'),
167
- gr.Dropdown(filter_choices, label="Filter", value=filter_choices[0])
168
- ],
169
- outputs=gr.Image(type="pil", label="Filtered Line Art"),
170
- title=title,
171
- description=description,
172
- examples=examples,
173
- allow_flagging='never'
174
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
175
 
176
  if __name__ == "__main__":
177
- iface.launch()
 
5
  from PIL import Image, ImageFilter, ImageOps, ImageChops
6
  import torchvision.transforms as transforms
7
  import os
8
+ import random
9
  import pathlib
10
 
11
  # --- ⚙️ Configuration ---
12
+ # Create a directory to save output images
13
  output_dir = "outputs"
14
  os.makedirs(output_dir, exist_ok=True)
15
 
16
+ # Define allowed image extensions for the file explorer
17
+ IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff"]
18
+
19
  # --- 🎨 Filters ---
20
  FILTERS = {
21
  "Standard": "📄", "Invert": "⚫⚪", "Blur": "🌫️", "Sharpen": "🔪", "Contour": "🗺️",
 
110
  return Image.fromarray(noisy_array)
111
  return line_img
112
 
113
+ # --- 🖼️ Main Processing Function (Updated) ---
114
+ def process_image(input_img_path, line_style, filter_choice, gallery_state):
115
  if not model1 or not model2:
116
  raise gr.Error("Models are not loaded. Please check for 'model.pth' and 'model2.pth'.")
117
+ if not input_img_path:
118
+ raise gr.Error("Please select an image from the file explorer first.")
119
+
120
  filter_name = filter_choice.split(" ", 1)[1]
121
  original_img = Image.open(input_img_path).convert('RGB')
122
 
 
135
 
136
  final_image = apply_filter(line_drawing_full_res, filter_name, original_img)
137
 
138
+ # --- 💾 Save the output and update gallery state ---
139
  base_name = pathlib.Path(input_img_path).stem
140
  output_filename = f"{base_name}_{filter_name}.png"
141
  output_filepath = os.path.join(output_dir, output_filename)
142
  final_image.save(output_filepath)
143
 
144
+ # Add new image path to the beginning of the list
145
+ gallery_state.insert(0, output_filepath)
146
+
147
+ # Return the single latest image for the main output and the updated list for the gallery
148
+ return final_image, gallery_state
149
 
150
  # --- 🚀 Gradio UI Setup ---
151
  title = "🖌️ Image to Line Art with Creative Filters"
152
+ description = "1. Browse and select an image using the file explorer. 2. Choose a line style. 3. Pick a filter. Your results will be saved to the 'outputs' folder and appear in the gallery below."
153
+
154
+ # --- Dynamic Examples Generation ---
155
+ def generate_examples():
156
+ example_images = [f"{i:02d}.jpeg" for i in range(1, 11)]
157
+ # Filter for only existing example images
158
+ valid_example_images = [img for img in example_images if os.path.exists(img)]
159
+
160
+ if not valid_example_images:
161
+ print("⚠️ Warning: No example images ('01.jpeg' through '10.jpeg') found. Examples will be empty.")
162
+ return []
163
+
164
+ examples = []
165
+ for name, emoji in FILTERS.items():
166
+ filter_choice = f"{emoji} {name}"
167
+ random_image = random.choice(valid_example_images)
168
+ line_style = random.choice(['Simple Lines', 'Complex Lines'])
169
+ examples.append([random_image, line_style, filter_choice])
170
+ # Shuffle to make the order random on each load
171
+ random.shuffle(examples)
172
+ return examples
173
+
174
+ with gr.Blocks(theme=gr.themes.Soft()) as demo:
175
+ gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
176
+ gr.Markdown(description)
177
+
178
+ # Stores the list of gallery image paths
179
+ gallery_state = gr.State(value=[])
180
+
181
+ with gr.Row():
182
+ with gr.Column(scale=1):
183
+ gr.Markdown("### 1. Select an Image")
184
+ # File explorer for a better user experience
185
+ input_image_path = gr.FileExplorer(
186
+ root=".",
187
+ glob=f"**/*[{''.join(ext[1:] for ext in IMAGE_EXTENSIONS)}]",
188
+ label="Browse Your Images",
189
+ height=400
190
+ )
191
+ gr.Markdown("### 2. Choose a Line Style")
192
+ line_style_radio = gr.Radio(
193
+ ['Complex Lines', 'Simple Lines'],
194
+ label="Line Style",
195
+ value='Simple Lines'
196
+ )
197
+
198
+ with gr.Column(scale=3):
199
+ gr.Markdown("### 3. Pick a Filter")
200
+ filter_buttons = [gr.Button(value=f"{emoji} {name}") for name, emoji in FILTERS.items()]
201
+
202
+ # Hidden radio to store the selected button's value
203
+ selected_filter = gr.Radio(
204
+ [b.value for b in filter_buttons],
205
+ label="Selected Filter",
206
+ visible=False,
207
+ value=filter_buttons[0].value
208
+ )
209
+
210
+ gr.Markdown("### 4. Result")
211
+ main_output_image = gr.Image(type="pil", label="Latest Result")
212
+
213
+ with gr.Row():
214
+ gr.Markdown("---")
215
+
216
+ with gr.Row():
217
+ # --- Dynamic Examples ---
218
+ gr.Examples(
219
+ examples=generate_examples(),
220
+ inputs=[input_image_path, line_style_radio, selected_filter],
221
+ label="✨ Click an Example to Start",
222
+ examples_per_page=10
223
+ )
224
+
225
+ with gr.Row():
226
+ gr.Markdown("## 🖼️ Result Gallery (Saved in 'outputs' folder)")
227
+ gallery_output = gr.Gallery(label="Your Generated Images", height=600, columns=5)
228
+
229
+ # --- Event Handling ---
230
+ def handle_filter_click(btn_value, current_path, style, state):
231
+ # When a filter button is clicked, it triggers the main processing function
232
+ new_main_img, new_state = process_image(current_path, style, btn_value, state)
233
+ # Update the hidden radio, the main image, and the gallery
234
+ return btn_value, new_main_img, new_state
235
+
236
+ for btn in filter_buttons:
237
+ btn.click(
238
+ fn=handle_filter_click,
239
+ inputs=[btn, input_image_path, line_style_radio, gallery_state],
240
+ outputs=[selected_filter, main_output_image, gallery_state]
241
+ )
242
 
243
  if __name__ == "__main__":
244
+ demo.launch()