@@ -12,7 +12,8 @@ def load_and_preprocess_image(image_path):
12
12
preprocess = transforms .Compose ([
13
13
transforms .Resize ((512 , 512 )),
14
14
transforms .ToTensor (),
15
- transforms .Normalize (mean = [0.485 , 0.456 , 0.406 ], std = [0.229 , 0.224 , 0.225 ]),
15
+ transforms .Normalize (mean = [0.485 , 0.456 , 0.406 ], std = [
16
+ 0.229 , 0.224 , 0.225 ]),
16
17
])
17
18
input_tensor = preprocess (image ).unsqueeze (0 )
18
19
return input_tensor
@@ -38,12 +39,15 @@ def remove_background(image_path, save_path, alpha_foreground=255, alpha_backgro
38
39
mask_pil = mask_pil .resize ((image .size [0 ], image .size [1 ]))
39
40
40
41
# Apply the mask to the input image
41
- image_with_alpha = Image .alpha_composite (Image .new ('RGBA' , image .size , (255 , 255 , 255 , alpha_background )), image )
42
- image_with_alpha .putalpha (mask_pil .point (lambda p : alpha_foreground if p else 0 ))
42
+ image_with_alpha = Image .alpha_composite (
43
+ Image .new ('RGBA' , image .size , (255 , 255 , 255 , alpha_background )), image )
44
+ image_with_alpha .putalpha (mask_pil .point (
45
+ lambda p : alpha_foreground if p else 0 ))
43
46
44
47
# Save the resulting image with transparent background
45
48
image_with_alpha .save (save_path , format = 'PNG' )
46
49
50
+
47
51
if __name__ == "__main__" :
48
52
# Replace 'input_folder' with the path to the folder containing your input images
49
53
input_folder = "/config/workspace/project/img"
@@ -54,5 +58,6 @@ def remove_background(image_path, save_path, alpha_foreground=255, alpha_backgro
54
58
for filename in os .listdir (input_folder ):
55
59
if filename .endswith (".jpg" ) or filename .endswith (".png" ) or filename .endswith (".avif" ):
56
60
image_path = os .path .join (input_folder , filename )
57
- save_path = os .path .join (output_folder , filename .replace (".jpg" , ".png" ).replace (".png" , ".png" ))
61
+ save_path = os .path .join (output_folder , filename .replace (
62
+ ".jpg" , ".png" ).replace (".png" , ".png" ))
58
63
remove_background (image_path , save_path )
0 commit comments