muhammadsalmanalfaridzi commited on
Commit
622f846
β€’
1 Parent(s): aea77c0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +699 -120
app.py CHANGED
@@ -1,142 +1,721 @@
 
 
 
 
 
 
 
1
  import gradio as gr
 
 
 
2
  import numpy as np
3
- import random
4
- #import spaces #[uncomment to use ZeroGPU]
5
- from diffusers import DiffusionPipeline
6
  import torch
7
 
8
- device = "cuda" if torch.cuda.is_available() else "cpu"
9
- model_repo_id = "stabilityai/sdxl-turbo" #Replace to the model you would like to use
 
 
10
 
11
- if torch.cuda.is_available():
12
- torch_dtype = torch.float16
13
- else:
14
- torch_dtype = torch.float32
15
 
16
- pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
17
- pipe = pipe.to(device)
 
 
 
18
 
19
- MAX_SEED = np.iinfo(np.int32).max
20
- MAX_IMAGE_SIZE = 1024
 
 
 
 
 
21
 
22
- #@spaces.GPU #[uncomment to use ZeroGPU]
23
- def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
 
 
 
24
 
25
- if randomize_seed:
26
- seed = random.randint(0, MAX_SEED)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
27
 
28
- generator = torch.Generator().manual_seed(seed)
29
-
30
- image = pipe(
31
- prompt = prompt,
32
- negative_prompt = negative_prompt,
33
- guidance_scale = guidance_scale,
34
- num_inference_steps = num_inference_steps,
35
- width = width,
36
- height = height,
37
- generator = generator
38
- ).images[0]
39
-
40
- return image, seed
41
-
42
- examples = [
43
- "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
44
- "An astronaut riding a green horse",
45
- "A delicious ceviche cheesecake slice",
46
- ]
47
-
48
- css="""
49
- #col-container {
50
- margin: 0 auto;
51
- max-width: 640px;
52
- }
53
- """
54
-
55
- with gr.Blocks(css=css) as demo:
56
-
57
- with gr.Column(elem_id="col-container"):
58
- gr.Markdown(f"""
59
- # Text-to-Image Gradio Template
60
- """)
61
 
62
- with gr.Row():
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
63
 
64
- prompt = gr.Text(
65
- label="Prompt",
66
- show_label=False,
67
- max_lines=1,
68
- placeholder="Enter your prompt",
69
- container=False,
70
- )
71
 
72
- run_button = gr.Button("Run", scale=0)
73
-
74
- result = gr.Image(label="Result", show_label=False)
75
-
76
- with gr.Accordion("Advanced Settings", open=False):
 
77
 
78
- negative_prompt = gr.Text(
79
- label="Negative prompt",
80
- max_lines=1,
81
- placeholder="Enter a negative prompt",
82
- visible=False,
83
- )
84
 
85
- seed = gr.Slider(
86
- label="Seed",
87
- minimum=0,
88
- maximum=MAX_SEED,
89
- step=1,
90
- value=0,
91
- )
 
 
 
 
 
 
 
 
92
 
93
- randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
 
 
 
 
 
 
 
94
 
95
- with gr.Row():
96
 
97
- width = gr.Slider(
98
- label="Width",
99
- minimum=256,
100
- maximum=MAX_IMAGE_SIZE,
101
- step=32,
102
- value=1024, #Replace with defaults that work for your model
103
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
104
 
105
- height = gr.Slider(
106
- label="Height",
107
- minimum=256,
108
- maximum=MAX_IMAGE_SIZE,
109
- step=32,
110
- value=1024, #Replace with defaults that work for your model
111
- )
112
-
113
- with gr.Row():
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
114
 
115
- guidance_scale = gr.Slider(
116
- label="Guidance scale",
117
- minimum=0.0,
118
- maximum=10.0,
119
- step=0.1,
120
- value=0.0, #Replace with defaults that work for your model
121
- )
122
 
123
- num_inference_steps = gr.Slider(
124
- label="Number of inference steps",
125
- minimum=1,
126
- maximum=50,
127
- step=1,
128
- value=2, #Replace with defaults that work for your model
129
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
130
 
131
- gr.Examples(
132
- examples = examples,
133
- inputs = [prompt]
134
- )
135
- gr.on(
136
- triggers=[run_button.click, prompt.submit],
137
- fn = infer,
138
- inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
139
- outputs = [result, seed]
140
- )
141
-
142
- demo.queue().launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import zipfile
3
+ import shutil
4
+ import time
5
+ from PIL import Image, ImageDraw
6
+ import io
7
+ from rembg import remove
8
  import gradio as gr
9
+ from concurrent.futures import ThreadPoolExecutor
10
+ from diffusers import StableDiffusionPipeline
11
+ from transformers import pipeline
12
  import numpy as np
13
+ import json
 
 
14
  import torch
15
 
16
+ # Load Stable Diffusion Model
17
+ def load_stable_diffusion_model():
18
+ device = "cuda" if torch.cuda.is_available() else "cpu"
19
+ return StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16).to(device)
20
 
21
+ sd_model = load_stable_diffusion_model()
 
 
 
22
 
23
+ def remove_background_stable_diffusion(image_path):
24
+ image = Image.open(image_path).convert("RGB")
25
+ prompt = "Remove background from the image"
26
+ output = sd_model(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
27
+ return output
28
 
29
+ def remove_background_rembg(input_path):
30
+ print(f"Removing background using rembg for image: {input_path}")
31
+ with open(input_path, 'rb') as i:
32
+ input_image = i.read()
33
+ output_image = remove(input_image)
34
+ img = Image.open(io.BytesIO(output_image)).convert("RGBA")
35
+ return img
36
 
37
+ def remove_background_bria(input_path):
38
+ print(f"Removing background using bria for image: {input_path}")
39
+ pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=0)
40
+ pillow_image = pipe(input_path)
41
+ return pillow_image
42
 
43
+ # Fungsi untuk memproses gambar menggunakan prompt
44
+ def text_to_image(prompt):
45
+ image = pipe(prompt).images[0]
46
+ image_path = f"generated_images/{prompt.replace(' ', '_')}.png"
47
+ image.save(image_path)
48
+ return image, image_path
49
+
50
+ def text_image_to_image(image, prompt):
51
+ # Convert input image to PIL Image for processing
52
+ modified_image = pipe(prompt, init_image=image, strength=0.75).images[0]
53
+ image_path = f"generated_images/{prompt.replace(' ', '_')}_modified.png"
54
+ modified_image.save(image_path)
55
+ return modified_image, image_path
56
+
57
+ def process_images(image_paths):
58
+ with ThreadPoolExecutor() as executor:
59
+ results = list(executor.map(remove_background_rembg, image_paths))
60
+ return results
61
+
62
+ def get_bounding_box_with_threshold(image, threshold):
63
+ # Convert image to numpy array
64
+ img_array = np.array(image)
65
+
66
+ # Get alpha channel
67
+ alpha = img_array[:,:,3]
68
+
69
+ # Find rows and columns where alpha > threshold
70
+ rows = np.any(alpha > threshold, axis=1)
71
+ cols = np.any(alpha > threshold, axis=0)
72
+
73
+ # Find the bounding box
74
+ top, bottom = np.where(rows)[0][[0, -1]]
75
+ left, right = np.where(cols)[0][[0, -1]]
76
+
77
+ if left < right and top < bottom:
78
+ return (left, top, right, bottom)
79
+ else:
80
+ return None
81
+
82
+ def check_cropped_sides(image, tolerance):
83
+ cropped_sides = []
84
+ width, height = image.size
85
+ edges = {
86
+ "top": [(x, 0) for x in range(width)],
87
+ "bottom": [(x, height - 1) for x in range(width)],
88
+ "left": [(0, y) for y in range(height)],
89
+ "right": [(width - 1, y) for y in range(height)]
90
+ }
91
+
92
+ for side, pixels in edges.items():
93
+ if any(image.getpixel(pixel)[3] > tolerance for pixel in pixels):
94
+ cropped_sides.append(side)
95
+
96
+ return cropped_sides
97
+
98
+ def resize_image(image, target_size, aspect_ratio):
99
+ target_width, target_height = target_size
100
+ if aspect_ratio > 1: # Landscape
101
+ new_height = target_height
102
+ new_width = int(new_height * aspect_ratio)
103
+ else: # Portrait or square
104
+ new_width = target_width
105
+ new_height = int(new_width / aspect_ratio)
106
+
107
+ return image.resize((new_width, new_height), Image.LANCZOS), new_width, new_height
108
+
109
+ def position_logic(image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left, use_threshold=True):
110
+ image = Image.open(image_path)
111
+ image = image.convert("RGBA")
112
+
113
+ # Get the bounding box of the non-blank area with threshold
114
+ if use_threshold:
115
+ bbox = get_bounding_box_with_threshold(image, threshold=10)
116
+ else:
117
+ bbox = image.getbbox()
118
+ log = []
119
+
120
+ if bbox:
121
+ # Check 1 pixel around the image for non-transparent pixels
122
+ width, height = image.size
123
+ cropped_sides = []
124
+
125
+ # Define tolerance for transparency
126
+ tolerance = 30 # Adjust this value as needed
127
+
128
+ # Check top edge
129
+ if any(image.getpixel((x, 0))[3] > tolerance for x in range(width)):
130
+ cropped_sides.append("top")
131
+
132
+ # Check bottom edge
133
+ if any(image.getpixel((x, height-1))[3] > tolerance for x in range(width)):
134
+ cropped_sides.append("bottom")
135
+
136
+ # Check left edge
137
+ if any(image.getpixel((0, y))[3] > tolerance for y in range(height)):
138
+ cropped_sides.append("left")
139
+
140
+ # Check right edge
141
+ if any(image.getpixel((width-1, y))[3] > tolerance for y in range(height)):
142
+ cropped_sides.append("right")
143
+
144
+ if cropped_sides:
145
+ info_message = f"Info for {os.path.basename(image_path)}: The following sides of the image may contain cropped objects: {', '.join(cropped_sides)}"
146
+ print(info_message)
147
+ log.append({"info": info_message})
148
+ else:
149
+ info_message = f"Info for {os.path.basename(image_path)}: The image is not cropped."
150
+ print(info_message)
151
+ log.append({"info": info_message})
152
+
153
+ # Crop the image to the bounding box
154
+ image = image.crop(bbox)
155
+ log.append({"action": "crop", "bbox": [str(bbox[0]), str(bbox[1]), str(bbox[2]), str(bbox[3])]})
156
+
157
+ # Calculate the new size to expand the image
158
+ target_width, target_height = canvas_size
159
+ aspect_ratio = image.width / image.height
160
+
161
+ if len(cropped_sides) == 4:
162
+ # If the image is cropped on all sides, center crop it to fit the canvas
163
+ if aspect_ratio > 1: # Landscape
164
+ new_height = target_height
165
+ new_width = int(new_height * aspect_ratio)
166
+ left = (new_width - target_width) // 2
167
+ image = image.resize((new_width, new_height), Image.LANCZOS)
168
+ image = image.crop((left, 0, left + target_width, target_height))
169
+ else: # Portrait or square
170
+ new_width = target_width
171
+ new_height = int(new_width / aspect_ratio)
172
+ top = (new_height - target_height) // 2
173
+ image = image.resize((new_width, new_height), Image.LANCZOS)
174
+ image = image.crop((0, top, target_width, top + target_height))
175
+ log.append({"action": "center_crop_resize", "new_size": f"{target_width}x{target_height}"})
176
+ x, y = 0, 0
177
+ elif not cropped_sides:
178
+ # If the image is not cropped, expand it from center until it touches the padding
179
+ new_height = target_height - padding_top - padding_bottom
180
+ new_width = int(new_height * aspect_ratio)
181
 
182
+ if new_width > target_width - padding_left - padding_right:
183
+ # If width exceeds available space, adjust based on width
184
+ new_width = target_width - padding_left - padding_right
185
+ new_height = int(new_width / aspect_ratio)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
186
 
187
+ # Resize the image
188
+ image = image.resize((new_width, new_height), Image.LANCZOS)
189
+ log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
190
+
191
+ x = (target_width - new_width) // 2
192
+ y = target_height - new_height - padding_bottom
193
+ else:
194
+ # New logic for handling cropped top and left, or top and right
195
+ if set(cropped_sides) == {"top", "left"} or set(cropped_sides) == {"top", "right"}:
196
+ new_height = target_height - padding_bottom
197
+ new_width = int(new_height * aspect_ratio)
198
+
199
+ # If new width exceeds canvas width, adjust based on width
200
+ if new_width > target_width:
201
+ new_width = target_width
202
+ new_height = int(new_width / aspect_ratio)
203
+
204
+ # Resize the image
205
+ image = image.resize((new_width, new_height), Image.LANCZOS)
206
+ log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
207
 
208
+ # Set position
209
+ if "left" in cropped_sides:
210
+ x = 0
211
+ else: # right in cropped_sides
212
+ x = target_width - new_width
213
+ y = 0
 
214
 
215
+ # If the resized image is taller than the canvas minus padding, crop from the bottom
216
+ if new_height > target_height - padding_bottom:
217
+ crop_bottom = new_height - (target_height - padding_bottom)
218
+ image = image.crop((0, 0, new_width, new_height - crop_bottom))
219
+ new_height = target_height - padding_bottom
220
+ log.append({"action": "crop_vertical", "bottom_pixels_removed": str(crop_bottom)})
221
 
222
+ log.append({"action": "position", "x": str(x), "y": str(y)})
223
+ elif set(cropped_sides) == {"bottom", "left"} or set(cropped_sides) == {"bottom", "right"}:
224
+ # Handle bottom & left or bottom & right cropped images
225
+ new_height = target_height - padding_top
226
+ new_width = int(new_height * aspect_ratio)
 
227
 
228
+ # If new width exceeds canvas width, adjust based on width
229
+ if new_width > target_width - padding_left - padding_right:
230
+ new_width = target_width - padding_left - padding_right
231
+ new_height = int(new_width / aspect_ratio)
232
+
233
+ # Resize the image without cropping or stretching
234
+ image = image.resize((new_width, new_height), Image.LANCZOS)
235
+ log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
236
+
237
+ # Set position
238
+ if "left" in cropped_sides:
239
+ x = 0
240
+ else: # right in cropped_sides
241
+ x = target_width - new_width
242
+ y = target_height - new_height
243
 
244
+ log.append({"action": "position", "x": str(x), "y": str(y)})
245
+ elif set(cropped_sides) == {"bottom", "left", "right"}:
246
+ # Expand the image from the center
247
+ new_width = target_width
248
+ new_height = int(new_width / aspect_ratio)
249
+
250
+ if new_height < target_height:
251
+ new_height = target_height
252
+ new_width = int(new_height * aspect_ratio)
253
 
254
+ image = image.resize((new_width, new_height), Image.LANCZOS)
255
 
256
+ # Crop to fit the canvas
257
+ left = (new_width - target_width) // 2
258
+ top = 0
259
+ image = image.crop((left, top, left + target_width, top + target_height))
260
+
261
+ log.append({"action": "expand_and_crop", "new_size": f"{target_width}x{target_height}"})
262
+ x, y = 0, 0
263
+ elif cropped_sides == ["top"]:
264
+ # New logic for handling only top-cropped images
265
+ if image.width > image.height:
266
+ new_width = target_width
267
+ new_height = int(target_width / aspect_ratio)
268
+ else:
269
+ new_height = target_height - padding_bottom
270
+ new_width = int(new_height * aspect_ratio)
271
+
272
+ # Resize the image
273
+ image = image.resize((new_width, new_height), Image.LANCZOS)
274
+ log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
275
+
276
+ x = (target_width - new_width) // 2
277
+ y = 0 # Align to top
278
+
279
+ # Apply padding only to non-cropped sides
280
+ x = max(padding_left, min(x, target_width - new_width - padding_right))
281
+ elif cropped_sides in [["right"], ["left"]]:
282
+ # New logic for handling only right-cropped or left-cropped images
283
+ if image.width > image.height:
284
+ new_width = target_width - max(padding_left, padding_right)
285
+ new_height = int(new_width / aspect_ratio)
286
+ else:
287
+ new_height = target_height - padding_top - padding_bottom
288
+ new_width = int(new_height * aspect_ratio)
289
 
290
+ # Resize the image
291
+ image = image.resize((new_width, new_height), Image.LANCZOS)
292
+ log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
293
+
294
+ if cropped_sides == ["right"]:
295
+ x = target_width - new_width # Align to right
296
+ else: # cropped_sides == ["left"]
297
+ x = 0 # Align to left
298
+ y = target_height - new_height - padding_bottom # Respect bottom padding
299
+
300
+ # Ensure top padding is respected
301
+ if y < padding_top:
302
+ y = padding_top
303
+
304
+ log.append({"action": "position", "x": str(x), "y": str(y)})
305
+ elif set(cropped_sides) == {"left", "right"}:
306
+ # Logic for handling images cropped on both left and right sides
307
+ new_width = target_width # Expand to full width of canvas
308
+
309
+ # Calculate the aspect ratio of the original image
310
+ aspect_ratio = image.width / image.height
311
+
312
+ # Calculate the new height while maintaining aspect ratio
313
+ new_height = int(new_width / aspect_ratio)
314
+
315
+ # Resize the image
316
+ image = image.resize((new_width, new_height), Image.LANCZOS)
317
+ log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
318
+
319
+ # Set horizontal position (always 0 as it spans full width)
320
+ x = 0
321
+
322
+ # Calculate vertical position to respect bottom padding
323
+ y = target_height - new_height - padding_bottom
324
+
325
+ # If the resized image is taller than the canvas, crop from the top only
326
+ if new_height > target_height - padding_bottom:
327
+ crop_top = new_height - (target_height - padding_bottom)
328
+ image = image.crop((0, crop_top, new_width, new_height))
329
+ new_height = target_height - padding_bottom
330
+ y = 0
331
+ log.append({"action": "crop_vertical", "top_pixels_removed": str(crop_top)})
332
+ else:
333
+ # Align the image to the bottom with padding
334
+ y = target_height - new_height - padding_bottom
335
+
336
+ log.append({"action": "position", "x": str(x), "y": str(y)})
337
+ elif cropped_sides == ["bottom"]:
338
+ # Logic for handling images cropped on the bottom side
339
+ # Calculate the aspect ratio of the original image
340
+ aspect_ratio = image.width / image.height
341
+
342
+ if aspect_ratio < 1: # Portrait orientation
343
+ new_height = target_height - padding_top # Full height with top padding
344
+ new_width = int(new_height * aspect_ratio)
345
 
346
+ # If the new width exceeds the canvas width, adjust it
347
+ if new_width > target_width:
348
+ new_width = target_width
349
+ new_height = int(new_width / aspect_ratio)
350
+ else: # Landscape orientation
351
+ new_width = target_width - padding_left - padding_right
352
+ new_height = int(new_width / aspect_ratio)
353
 
354
+ # If the new height exceeds the canvas height, adjust it
355
+ if new_height > target_height:
356
+ new_height = target_height
357
+ new_width = int(new_height * aspect_ratio)
358
+
359
+ # Resize the image
360
+ image = image.resize((new_width, new_height), Image.LANCZOS)
361
+ log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
362
+
363
+ # Set horizontal position (centered)
364
+ x = (target_width - new_width) // 2
365
+
366
+ # Set vertical position (touching bottom edge for all cases)
367
+ y = target_height - new_height
368
+
369
+ log.append({"action": "position", "x": str(x), "y": str(y)})
370
+ else:
371
+ # Use the original resizing logic for other partially cropped images
372
+ if image.width > image.height:
373
+ new_width = target_width
374
+ new_height = int(target_width / aspect_ratio)
375
+ else:
376
+ new_height = target_height
377
+ new_width = int(target_height * aspect_ratio)
378
+
379
+ # Resize the image
380
+ image = image.resize((new_width, new_height), Image.LANCZOS)
381
+ log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
382
+
383
+ # Center horizontally for all images
384
+ x = (target_width - new_width) // 2
385
+ y = target_height - new_height - padding_bottom
386
+
387
+ # Adjust positions for cropped sides
388
+ if "top" in cropped_sides:
389
+ y = 0
390
+ elif "bottom" in cropped_sides:
391
+ y = target_height - new_height
392
+ if "left" in cropped_sides:
393
+ x = 0
394
+ elif "right" in cropped_sides:
395
+ x = target_width - new_width
396
+
397
+ # Apply padding only to non-cropped sides, but keep horizontal centering
398
+ if "left" not in cropped_sides and "right" not in cropped_sides:
399
+ x = (target_width - new_width) // 2 # Always center horizontally
400
+ if "top" not in cropped_sides and "bottom" not in cropped_sides:
401
+ y = max(padding_top, min(y, target_height - new_height - padding_bottom))
402
+
403
+ return log, image, x, y
404
+
405
+ def get_canvas_size(canvas_size_name):
406
+ sizes = {
407
+ 'Rox': ((1080, 1080), (112, 125, 116, 125)),
408
+ 'Columbia': ((730, 610), (30, 105, 35, 105)),
409
+ 'Zalora': ((763, 1100), (50, 50, 200, 50)),
410
+ }
411
+ return sizes.get(canvas_size_name, ((1080, 1080), (0, 0, 0, 0)))
412
+
413
+ def process_single_image(image_path, output_folder, bg_method, canvas_size_name, output_format, bg_choice, custom_color, watermark_path=None):
414
+ add_padding_line = False
415
+
416
+ if canvas_size_name == 'Rox':
417
+ canvas_size = (1080, 1080)
418
+ padding_top = 112
419
+ padding_right = 125
420
+ padding_bottom = 116
421
+ padding_left = 125
422
+ elif canvas_size_name == 'Columbia':
423
+ canvas_size = (730, 610)
424
+ padding_top = 30
425
+ padding_right = 105
426
+ padding_bottom = 35
427
+ padding_left = 105
428
+ elif canvas_size_name == 'Zalora':
429
+ canvas_size = (763, 1100)
430
+ padding_top = 50
431
+ padding_right = 50
432
+ padding_bottom = 200
433
+ padding_left = 50
434
+
435
+ filename = os.path.basename(image_path)
436
+ try:
437
+ print(f"Processing image: {filename}")
438
+ if bg_method == 'stable_diffusion':
439
+ image_with_no_bg = remove_background_stable_diffusion(image_path)
440
+ elif bg_method == 'rembg':
441
+ image_with_no_bg = remove_background_rembg(image_path) # Placeholder for existing function
442
+ elif bg_method == 'bria':
443
+ image_with_no_bg = remove_background_bria(image_path) # Placeholder for existing function
444
 
445
+ temp_image_path = os.path.join(output_folder, f"temp_{filename}")
446
+ image_with_no_bg.save(temp_image_path, format='PNG')
447
+
448
+ log, new_image, x, y = position_logic(temp_image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left)
449
+
450
+ # Create a new canvas with the appropriate background
451
+ if bg_choice == 'white':
452
+ canvas = Image.new("RGBA", canvas_size, "WHITE")
453
+ elif bg_choice == 'custom':
454
+ canvas = Image.new("RGBA", canvas_size, custom_color)
455
+ else: # transparent
456
+ canvas = Image.new("RGBA", canvas_size, (0, 0, 0, 0))
457
+
458
+ # Paste the resized image onto the canvas
459
+ canvas.paste(new_image, (x, y), new_image)
460
+ log.append({"action": "paste", "position": [str(x), str(y)]})
461
+
462
+ # Add visible black line for padding when background is not transparent
463
+ if add_padding_line:
464
+ draw = ImageDraw.Draw(canvas)
465
+ draw.rectangle([padding_left, padding_top, canvas_size[0] - padding_right, canvas_size[1] - padding_bottom], outline="black", width=5)
466
+ log.append({"action": "add_padding_line"})
467
+
468
+ output_ext = 'jpg' if output_format == 'JPG' else 'png'
469
+ output_filename = f"{os.path.splitext(filename)[0]}.{output_ext}"
470
+ output_path = os.path.join(output_folder, output_filename)
471
+
472
+ # Apply watermark only if the filename ends with "_01" and watermark_path is provided
473
+ if os.path.splitext(filename)[0].endswith("_01") and watermark_path:
474
+ watermark = Image.open(watermark_path).convert("RGBA")
475
+ canvas.paste(watermark, (0, 0), watermark)
476
+ log.append({"action": "add_watermark"})
477
+
478
+ if output_format == 'JPG':
479
+ canvas.convert('RGB').save(output_path, format='JPEG')
480
+ else:
481
+ canvas.save(output_path, format='PNG')
482
+
483
+ os.remove(temp_image_path)
484
+
485
+ print(f"Processed image path: {output_path}")
486
+ return [(output_path, image_path)], log
487
+
488
+ except Exception as e:
489
+ print(f"Error processing {filename}: {e}")
490
+ return None, None
491
+
492
+ def create_canvas(canvas_size, bg_choice, custom_color):
493
+ if bg_choice == 'white':
494
+ return Image.new("RGBA", canvas_size, "WHITE")
495
+ elif bg_choice == 'custom':
496
+ return Image.new("RGBA", canvas_size, custom_color)
497
+ else: # transparent
498
+ return Image.new("RGBA", canvas_size, (0, 0, 0, 0))
499
+
500
+ def apply_watermark(canvas, watermark_path):
501
+ watermark = Image.open(watermark_path).convert("RGBA")
502
+ canvas.paste(watermark, (0, 0), watermark)
503
+
504
+ def save_image(canvas, output_folder, filename, output_format):
505
+ output_ext = 'jpg' if output_format == 'JPG' else 'png'
506
+ output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.{output_ext}")
507
+ if output_format == 'JPG':
508
+ canvas.convert('RGB').save(output_path, format='JPEG')
509
+ else:
510
+ canvas.save(output_path, format='PNG')
511
+ return output_path
512
+
513
+ def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_size='Rox', output_format='PNG', bg_choice='transparent', custom_color="#ffffff", num_workers=4, progress=gr.Progress()):
514
+ start_time = time.time()
515
+
516
+ output_folder = "processed_images"
517
+ if os.path.exists(output_folder):
518
+ shutil.rmtree(output_folder)
519
+ os.makedirs(output_folder)
520
+
521
+ processed_images = []
522
+ original_images = []
523
+ all_logs = []
524
+
525
+ if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
526
+ input_folder = "temp_input"
527
+ if os.path.exists(input_folder):
528
+ shutil.rmtree(input_folder)
529
+ os.makedirs(input_folder)
530
+
531
+ try:
532
+ with zipfile.ZipFile(input_files, 'r') as zip_ref:
533
+ zip_ref.extractall(input_folder)
534
+ except zipfile.BadZipFile as e:
535
+ print(f"Error extracting zip file: {e}")
536
+ return [], None, 0
537
+
538
+ image_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp'))]
539
+ elif isinstance(input_files, list):
540
+ image_files = input_files
541
+ else:
542
+ image_files = [input_files]
543
+
544
+ total_images = len(image_files)
545
+ print(f"Total images to process: {total_images}")
546
+
547
+ avg_processing_time = 0
548
+ with ThreadPoolExecutor(max_workers=num_workers) as executor:
549
+ future_to_image = {executor.submit(process_single_image, image_path, output_folder, bg_method, canvas_size, output_format, bg_choice, custom_color, watermark_path): image_path for image_path in image_files}
550
+ for idx, future in enumerate(future_to_image):
551
+ try:
552
+ start_time_image = time.time()
553
+ result, log = future.result()
554
+ end_time_image = time.time()
555
+ image_processing_time = end_time_image - start_time_image
556
+
557
+ # Update average processing time
558
+ avg_processing_time = (avg_processing_time * idx + image_processing_time) / (idx + 1)
559
+
560
+ if result:
561
+ processed_images.extend(result)
562
+ original_images.append(future_to_image[future])
563
+ all_logs.append({os.path.basename(future_to_image[future]): log})
564
+
565
+ # Estimate remaining time
566
+ remaining_images = total_images - (idx + 1)
567
+ estimated_remaining_time = remaining_images * avg_processing_time
568
+
569
+ progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed. Estimated time remaining: {estimated_remaining_time:.2f} seconds")
570
+ except Exception as e:
571
+ print(f"Error processing image {future_to_image[future]}: {e}")
572
+
573
+ output_zip_path = "processed_images.zip"
574
+ with zipfile.ZipFile(output_zip_path, 'w') as zipf:
575
+ for file, _ in processed_images:
576
+ zipf.write(file, os.path.basename(file))
577
+
578
+ # Write the comprehensive log for all images
579
+ with open(os.path.join(output_folder, 'process_log.json'), 'w') as log_file:
580
+ json.dump(all_logs, log_file, indent=4)
581
+ print("Comprehensive log saved to", os.path.join(output_folder, 'process_log.json'))
582
+
583
+ end_time = time.time()
584
+ processing_time = end_time - start_time
585
+ print(f"Processing time: {processing_time} seconds")
586
+
587
+ return original_images, processed_images, output_zip_path, processing_time
588
+
589
+ def extract_image_files(input_files):
590
+ if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
591
+ input_folder = "temp_input"
592
+ if os.path.exists(input_folder):
593
+ shutil.rmtree(input_folder)
594
+ os.makedirs(input_folder)
595
+
596
+ with zipfile.ZipFile(input_files, 'r') as zip_ref:
597
+ zip_ref.extractall(input_folder)
598
+
599
+ return [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp'))]
600
+ elif isinstance(input_files, list):
601
+ return input_files
602
+ else:
603
+ return [input_files]
604
+
605
+ def create_output_zip(processed_images):
606
+ output_zip_path = "processed_images.zip"
607
+ with zipfile.ZipFile(output_zip_path, 'w') as zipf:
608
+ for file, _ in processed_images:
609
+ zipf.write(file, os.path.basename(file))
610
+ return output_zip_path
611
+
612
+ def save_log(all_logs, output_folder):
613
+ with open(os.path.join(output_folder, 'process_log.json'), 'w') as log_file:
614
+ json.dump(all_logs, log_file, indent=4)
615
+ print("Comprehensive log saved to", os.path.join(output_folder, 'process_log.json'))
616
+
617
+ def gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
618
+ progress = gr.Progress()
619
+ watermark_path = watermark.name if watermark else None
620
+
621
+ if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
622
+ return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
623
+ elif isinstance(input_files, list):
624
+ return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
625
+ else:
626
+ return process_images(input_files.name, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
627
+
628
+ def show_color_picker(bg_choice):
629
+ if bg_choice == 'custom':
630
+ return gr.update(visible=True)
631
+ return gr.update(visible=False)
632
+
633
+ def update_compare(evt: gr.SelectData):
634
+ if isinstance(evt.value, dict) and 'caption' in evt.value:
635
+ input_path = evt.value['caption']
636
+ output_path = evt.value['image']['path']
637
+ input_path = input_path.split("Input: ")[-1]
638
+
639
+ original_img = Image.open(input_path)
640
+ processed_img = Image.open(output_path)
641
+
642
+ original_ratio = f"{original_img.width}x{original_img.height}"
643
+ processed_ratio = f"{processed_img.width}x{processed_img.height}"
644
+
645
+ return gr.update(value=input_path), gr.update(value=output_path), gr.update(value=original_ratio), gr.update(value=processed_ratio)
646
+ else:
647
+ print("No caption found in selection")
648
+ return gr.update(value=None), gr.update(value=None), gr.update(value=None), gr.update(value=None)
649
+
650
+ def process(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
651
+ _, processed_images, zip_path, time_taken = gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers)
652
+ processed_images_with_captions = [(img, f"Input: {caption}") for img, caption in processed_images]
653
+ return processed_images_with_captions, zip_path, f"{time_taken:.2f} seconds"
654
+
655
+ with gr.Blocks(theme="NoCrypt/miku@1.2.2") as iface:
656
+ gr.Markdown("# Image Background Removal and Resizing with Optional Watermark")
657
+ gr.Markdown("Choose to upload multiple images or a ZIP/RAR file, select the crop mode, optionally upload a watermark image, and choose the output format.")
658
+
659
+ # Fitur Text to Image
660
+ gr.Markdown("# Text to Image Feature")
661
+ with gr.Row():
662
+ prompt_input = gr.Textbox(label="Enter your prompt for image generation:")
663
+ generate_button = gr.Button("Generate Image")
664
+ output_image = gr.Image(label="Generated Image")
665
+ download_button = gr.File(label="Download Generated Image", type="filepath")
666
+
667
+ generate_button.click(text_to_image, inputs=prompt_input, outputs=[output_image, download_button])
668
+
669
+ # Fitur Text Image to Image
670
+ gr.Markdown("# Text Image to Image Feature")
671
+ with gr.Row():
672
+ input_image = gr.Image(label="Upload Image for Modification", type="pil")
673
+ prompt_modification = gr.Textbox(label="Enter your prompt for modification:")
674
+ modify_button = gr.Button("Modify Image")
675
+ modified_output_image = gr.Image(label="Modified Image")
676
+ download_modified_button = gr.File(label="Download Modified Image", type="filepath")
677
+
678
+ modify_button.click(text_image_to_image, inputs=[input_image, prompt_modification], outputs=[modified_output_image, download_modified_button])
679
+
680
+ with gr.Row():
681
+ input_files = gr.File(label="Upload Image or ZIP/RAR file", file_types=[".zip", ".rar", "image"], interactive=True)
682
+ watermark = gr.File(label="Upload Watermark Image (Optional)", file_types=[".png"])
683
+
684
+ with gr.Row():
685
+ canvas_size = gr.Radio(choices=["Rox", "Columbia", "Zalora"], label="Canvas Size", value="Rox")
686
+ output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="JPG")
687
+ num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Number of Workers", value=5)
688
+
689
+ with gr.Row():
690
+ bg_method = gr.Radio(choices=["bria", "rembg", "stable_diffusion"], label="Background Removal Method", value="stable_diffusion")
691
+ bg_choice = gr.Radio(choices=["transparent", "white", "custom"], label="Background Choice", value="white")
692
+ custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False)
693
+
694
+ process_button = gr.Button("Process Images")
695
+
696
+ with gr.Row():
697
+ input_image = gr.File(label="Upload Image", file_types=["image"])
698
+ prompt = gr.Textbox(label="Prompt for Image Modification")
699
+ process_button = gr.Button("Generate Image")
700
+ output_image = gr.Image(label="Generated Image")
701
+
702
+ process_button.click(gradio_interface, inputs=[input_image, prompt], outputs=[output_image])
703
+
704
+
705
+ with gr.Row():
706
+ gallery_processed = gr.Gallery(label="Processed Images")
707
+ with gr.Row():
708
+ image_original = gr.Image(label="Original Images", interactive=False)
709
+ image_processed = gr.Image(label="Processed Images", interactive=False)
710
+ with gr.Row():
711
+ original_ratio = gr.Textbox(label="Original Ratio")
712
+ processed_ratio = gr.Textbox(label="Processed Ratio")
713
+ with gr.Row():
714
+ output_zip = gr.File(label="Download Processed Images as ZIP")
715
+ processing_time = gr.Textbox(label="Processing Time (seconds)")
716
+
717
+ bg_choice.change(show_color_picker, inputs=bg_choice, outputs=custom_color)
718
+ process_button.click(process, inputs=[input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers], outputs=[gallery_processed, output_zip, processing_time])
719
+ gallery_processed.select(update_compare, outputs=[image_original, image_processed, original_ratio, processed_ratio])
720
+
721
+ iface.launch(share=True)