File size: 2,810 Bytes
9754b95 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import os
from PIL import Image
from retinaface import RetinaFace
from concurrent.futures import ProcessPoolExecutor
Image.MAX_IMAGE_PIXELS = 200000000 # Allow images up to 200 million pixels
DATA = [
((512, 512), r"G:\man dataset\man_4502_imgs_raw"),
((1024, 1024), r"G:\man dataset\man_4502_imgs_raw")
]
NUM_THREADS = 1 # Modify this for the desired number of threads
OUTPUT_ROOT = r"H:\man dataset\resize_faces" # Root folder to save resized face images
## change resize enabled to resize
def resize_face(args, resize_enabled=False):
try:
input_path, resolution, output_root = args
except ValueError as e:
print(f"Error unpacking arguments: {e}")
return
try:
image = Image.open(input_path)
except Exception as e:
print(f"Error loading image: {e}")
return
try:
# Detect faces
faces = RetinaFace.detect_faces(input_path)
if not faces:
print(f"No face detected in {input_path}")
return
# Assuming we take the first detected face
face = next(iter(faces.values()))
x, y, x2, y2 = face['facial_area']
face_image = image.crop((x, y, x2, y2))
if resize_enabled:
# Resize the cropped face only if resizing is enabled
face_image = face_image.resize(resolution, Image.LANCZOS)
except Exception as e:
print(f"Error in face detection or resizing: {e}")
return
try:
output_folder = os.path.join(output_root, f"{resolution[0]}x{resolution[1]}" if resize_enabled else "no_resize")
if not os.path.exists(output_folder):
os.makedirs(output_folder)
output_path = os.path.join(output_folder, os.path.basename(input_path))
face_image.save(output_path, quality=100)
except Exception as e:
print(f"Error saving image: {e}")
def process_folder(input_folder, resolution, resize_enabled=False):
image_paths = [os.path.join(input_folder, fname) for fname in os.listdir(input_folder) if fname.lower().endswith(('png', 'jpg', 'jpeg'))]
total_images = len(image_paths)
processed_count = 0
with ProcessPoolExecutor(max_workers=NUM_THREADS) as executor:
for _ in executor.map(resize_face, [(path, resolution, OUTPUT_ROOT) for path in image_paths], [resize_enabled] * len(image_paths)):
processed_count += 1
print(f"Processed {processed_count}/{total_images} images for resolution {resolution[0]}x{resolution[1]} (Resize: {'Enabled' if resize_enabled else 'Disabled'})...")
if __name__ == "__main__":
resize_enabled = False # Change this to True if resizing is needed
for resolution, folder in DATA:
process_folder(folder, resolution, resize_enabled)
print("Processing complete!")
|