Script pré traitement images

Pour utiliser des logiciels de photogrammétrie il est important d’avoir des images uniformes et bien traitées.

Pour cela, il est souvent conseiller d’effacer le fond des images pour ne conserver que le sujet principal afin de faciliter le travail du logiciels.

Des outils comme le traitements par lot de photoshop existe mais est assez lourd, et donne des résultats trop variables.

J’ai donc préparé un script Python (avec l’aide de ChatGPT…) qui prends toutes les images d’un dossier et de ses sous dossiers, utilise une librairie open source « Rembg » qui identifie à l’aide d’algorithmes d’IA le sujet principal de l’image, et enlève le fond d’écran pour le remplacer par la couleur de notre choix.

Il est possible de choisir le modèle utilisé selon le type d’éléments qu’on détoure.

Enfin, il transfère les métadonnées du fichier d’origine (pour une meilleure gestion sur des logiciels comme Lightroom ou DigiKam) au fichiers créés ainsi que la structure des sous dossiers.

Voici le script :

import os
import cv2
import numpy as np
from rembg.bg import remove, new_session
from pyexiv2 import Image
# Replace the transparent background with a light gray color
# (255, 255, 255) => White // (0, 0, 0) => Black
background_color = (70, 200, 255)
# Define the input and output folder paths
# If input == output => remplace with edited file
# If subfolder in input, will reproduce on output
input_folder = "./input"
output_folder = "./output2"
#create a new session by passing name of the model from one of the following
#["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg", "silueta"]
#infos and difference on models : https://github.com/danielgatis/rembg
my_session = new_session("u2net_human_seg")
def add_background(image, background_color):
# Check if image has alpha channel
if image.shape[2] == 4:
# Extract alpha channel and convert to float between 0 and 1
alpha_channel = image[:, :, 3] / 255.0
# Convert image to RGB without alpha channel
image = image[:, :, :3]
# Create background with same shape as image
background = np.zeros_like(image)
# Set background color
background_color = np.array(background_color).reshape((1, 1, 3))
background[:, :] = background_color
# Blend image and background using alpha channel
blended = alpha_channel[..., np.newaxis] * image + (1 - alpha_channel[..., np.newaxis]) * background
# Convert blended image to uint8
blended = blended.astype(np.uint8)
return blended
else:
# Image doesn't have alpha channel, so just return it as is
return image
# Check if the output folder is the same as the input folder
if os.path.abspath(input_folder) == os.path.abspath(output_folder):
overwrite = input("The output folder is the same as the input folder. Are you sure you want to overwrite the source files? (y/n): ")
if overwrite.lower() != "y":
exit()
# Loop through all the JPEG files in the folder and its subfolders
for root, dirs, files in os.walk(input_folder):
for file in files:
if file.lower().endswith(".jpg"):
# Read the JPEG file as a binary stream
with open(os.path.join(root, file), "rb") as f:
img_data = f.read()
xmp_data = Image(os.path.join(root, file)).read_xmp()
# Use the remove() function from the rembg library to remove the background
output = remove(img_data, session=my_session)
# Convert the binary output to a NumPy array
np_img = cv2.imdecode(np.frombuffer(output, np.uint8), cv2.IMREAD_UNCHANGED)
# Create output folder if it doesn't exist
output_subfolder_path = os.path.join(output_folder, os.path.relpath(root, input_folder))
os.makedirs(output_subfolder_path, exist_ok=True)
result = add_background(np_img, background_color)
# Get the output file path
output_file_path = os.path.join(output_subfolder_path, file)
# Write the processed image to disk, overwriting the original file
cv2.imwrite(os.path.join(output_file_path), result)
print(output_file_path)
# Inject the stored XMP metadata into the processed file
output_image = Image(os.path.join(output_file_path))
output_image.modify_xmp(xmp_data)
output_image.close()
print(f"{file} processed.")
print("All images processed.")