@@ -12,7 +12,8 @@ def load_and_preprocess_image(image_path):
1212 preprocess = transforms .Compose ([
1313 transforms .Resize ((512 , 512 )),
1414 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 ]),
1617 ])
1718 input_tensor = preprocess (image ).unsqueeze (0 )
1819 return input_tensor
@@ -38,12 +39,15 @@ def remove_background(image_path, save_path, alpha_foreground=255, alpha_backgro
3839 mask_pil = mask_pil .resize ((image .size [0 ], image .size [1 ]))
3940
4041 # 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 ))
4346
4447 # Save the resulting image with transparent background
4548 image_with_alpha .save (save_path , format = 'PNG' )
4649
50+ 4751if __name__ == "__main__" :
4852 # Replace 'input_folder' with the path to the folder containing your input images
4953 input_folder = "/config/workspace/project/img"
@@ -54,5 +58,6 @@ def remove_background(image_path, save_path, alpha_foreground=255, alpha_backgro
5458 for filename in os .listdir (input_folder ):
5559 if filename .endswith (".jpg" ) or filename .endswith (".png" ) or filename .endswith (".avif" ):
5660 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" ))
5863 remove_background (image_path , save_path )
0 commit comments