First Commit
This commit is contained in:
commit
e297e484ee
4 changed files with 904 additions and 0 deletions
207
face_encoder_gpu.py
Normal file
207
face_encoder_gpu.py
Normal file
|
|
@ -0,0 +1,207 @@
|
|||
import sys
|
||||
import signal
|
||||
|
||||
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
||||
|
||||
import numpy as np
|
||||
import dlib
|
||||
import os
|
||||
import face_recognition
|
||||
import argparse
|
||||
import pickle
|
||||
from tqdm import tqdm
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
|
||||
date_time = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
default_log_filename = f"encoder_log_{date_time}.txt"
|
||||
default_out_filename = f"face_encodings_{date_time}.pkl"
|
||||
|
||||
def format_time(seconds, total_images):
|
||||
hours, rem = divmod(seconds, 3600)
|
||||
minutes, seconds_final = divmod(rem, 60)
|
||||
|
||||
time_str = ""
|
||||
if hours > 0: time_str += f"{int(hours)}h "
|
||||
if minutes > 0: time_str += f"{int(minutes)}m "
|
||||
time_str += f"{seconds_final:.2f}s"
|
||||
|
||||
avg_speed = total_images / seconds
|
||||
|
||||
return time_str, avg_speed
|
||||
|
||||
def resolve_path(path, default):
|
||||
default_dirname = "output"
|
||||
default_filename = default
|
||||
|
||||
if not path:
|
||||
resolved_path = Path(default_dirname) / default_filename
|
||||
return resolved_path.resolve()
|
||||
|
||||
resolved_path = Path(path).resolve()
|
||||
|
||||
if resolved_path.is_dir() or path.endswith(os.sep) or path.endswith('/') or not resolved_path.suffix:
|
||||
resolved_path = resolved_path / default_filename
|
||||
return resolved_path.resolve()
|
||||
|
||||
return resolved_path
|
||||
|
||||
def encode_images(images_dir, log, recursive=False, include_tn=False):
|
||||
encodings = []
|
||||
filenames = []
|
||||
|
||||
log_path = resolve_path(log, default_log_filename)
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if not dlib.DLIB_USE_CUDA:
|
||||
print("\n" + "#" * 80)
|
||||
print("ERRORE CRITICO: GPU non rilevata.")
|
||||
print("Il programma è configurato per funzionare esclusivamente con CUDA.")
|
||||
print("#" * 80 + "\n")
|
||||
with open(log_path, "w", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [ERRORE] GPU non rilevata ---\n")
|
||||
sys.exit(1)
|
||||
|
||||
model_type = "cnn"
|
||||
print("Modalità GPU (CUDA) rilevata e attivata correttamente.")
|
||||
with open(log_path, "w", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [INFO] Modalità GPU (CUDA) rilevata e attivata correttamente ---\n")
|
||||
|
||||
images_dir_path = Path(images_dir).resolve()
|
||||
if not images_dir_path.exists() or not images_dir_path.is_dir():
|
||||
print(f"Errore: La cartella {images_dir_path} non esiste.")
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [ERRORE] La cartella {images_dir_path} non esiste ---\n")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.JPG', '*.JPEG', '*.PNG', '*.BMP']
|
||||
files_to_process = []
|
||||
|
||||
if recursive:
|
||||
for ext in extensions:
|
||||
files_to_process.extend(images_dir_path.rglob(ext))
|
||||
else:
|
||||
for ext in extensions:
|
||||
files_to_process.extend(images_dir_path.glob(ext))
|
||||
|
||||
files_to_process = sorted(list(set(files_to_process)))
|
||||
|
||||
if not files_to_process:
|
||||
print("Nessuna immagine trovata da elaborare.")
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [INFO] Nessuna immagine trovata da elaborare ---\n")
|
||||
sys.exit(1)
|
||||
print(f"Trovate {len(files_to_process)} immagini da elaborare.")
|
||||
with open(log_path, "w", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [INFO] Trovate {len(files_to_process)} immagini da elaborare ---\n")
|
||||
|
||||
if not include_tn:
|
||||
total_images = len(files_to_process)
|
||||
files_to_process = [f for f in files_to_process if not f.name.lower().startswith("tn_")]
|
||||
print(f"Filtro-tn attivo. Rimosse {total_images - len(files_to_process)} immagini thumbnail. Rimaste {len(files_to_process)} immagini.")
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [INFO] Filtro-tn attivo. Rimosse {total_images - len(files_to_process)} immagini thumbnail. Rimaste {len(files_to_process)} immagini ---\n")
|
||||
else:
|
||||
print(f"Filtro-tn disattivato.")
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [INFO] Filtro-tn disattivato ---\n")
|
||||
|
||||
print(f"Avvio codifica immagini da {images_dir_path}{' in modalità ricorsiva' if recursive else ''})")
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [INFO] Codifica avviata da {images_dir_path} {'in modalità ricorsiva' if recursive else ''} ---\n")
|
||||
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"\n============== [INIZIO ELABORAZIONE] ==============\n\n")
|
||||
|
||||
dummy_image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
_ = face_recognition.face_locations(dummy_image, model=model_type)
|
||||
pbar = tqdm(total=len(files_to_process), desc="Elaborazione immagini", unit="img")
|
||||
start_time = datetime.now().timestamp()
|
||||
|
||||
signal.signal(signal.SIGINT, signal.default_int_handler)
|
||||
|
||||
try:
|
||||
for path in files_to_process:
|
||||
try:
|
||||
image = face_recognition.load_image_file(path)
|
||||
|
||||
face_locations = face_recognition.face_locations(image, model=model_type)
|
||||
face_encodings = face_recognition.face_encodings(image, known_face_locations=face_locations)
|
||||
|
||||
for face_encoding in face_encodings:
|
||||
encodings.append(face_encoding)
|
||||
filenames.append(str(path.relative_to(images_dir_path)))
|
||||
|
||||
nfaces = len(face_locations)
|
||||
msg = f"{path.relative_to(images_dir_path)} - [{nfaces:<2} {'volto' if nfaces == 1 else 'volti'}]"
|
||||
pbar.write(msg)
|
||||
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"{msg}\n")
|
||||
|
||||
except Exception as e:
|
||||
err_msg = f"Errore durante l'elaborazione di {path.name}: {e}"
|
||||
pbar.write(err_msg)
|
||||
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"--- [ERRORE] {err_msg} ---\n")
|
||||
|
||||
pbar.update(1)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
pbar.disable = True
|
||||
pbar.close()
|
||||
|
||||
print("\nInterruzione manuale dell'elaborazione, salvataggio dei dati finora elaborati...")
|
||||
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write("\n============== [ELABORAZIONE INTERROTTA MANUALMENTE] ==============\n")
|
||||
|
||||
else:
|
||||
pbar.close()
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write("\n============== [ELABORAZIONE COMPLETATA CON SUCCESSO] ==============\n")
|
||||
|
||||
finally:
|
||||
execution_time = datetime.now().timestamp() - start_time
|
||||
time_str, avg_speed = format_time(execution_time, len(set(filenames)))
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"\n--- [INFO] Tempo impiegato: {time_str} ---")
|
||||
log_f.write(f"\n--- [INFO] Velocità media: {avg_speed:.1f} img/s ---")
|
||||
|
||||
return encodings, filenames
|
||||
|
||||
def save_encodings(encodings, filenames, output, log):
|
||||
data = {"encodings": encodings, "filenames": filenames}
|
||||
output_path = resolve_path(output, default_out_filename)
|
||||
log_path = resolve_path(log, default_log_filename)
|
||||
|
||||
try:
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(output_path, "wb") as f:
|
||||
pickle.dump(data, f)
|
||||
print(f"Codifica terminata, encodings salvati in {output_path}")
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"\n--- [INFO] Codifica terminata, encodings salvati in {output_path} ---\n")
|
||||
except Exception as e:
|
||||
print(f"Errore di salvataggio: {e}")
|
||||
with open(log_path, "a", encoding="utf-8") as log_f:
|
||||
log_f.write(f"\n--- [ERRORE] Errore durante il salvataggio: {e} ---\n")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="VERSIONE GPU [CUDA].\nGenera gli encoding, codificando le foto 'unknown'.")
|
||||
parser.add_argument("-i", "--images", required=True, help="Cartella contenente le foto da codificare")
|
||||
parser.add_argument("-o", "--out", help="Percorso del file di output contentente gli encoding. Default: './output/face_encodings_[datetime].pkl'")
|
||||
parser.add_argument("-l", "--log", help="Percorso del file di log. Default: './output/encoder_log_[datetime].txt'")
|
||||
parser.add_argument("-r", "--recursive", action="store_true", help="Cerca immagini anche nelle sottocartelle")
|
||||
parser.add_argument("-t", "--include-tn", action="store_true", help="Include nell'elabortazione anche le immagini thumbnail che iniziano con 'tn_'")
|
||||
args = parser.parse_args()
|
||||
|
||||
encodings, filenames = encode_images(args.images, args.log, args.recursive, args.include_tn)
|
||||
|
||||
if encodings:
|
||||
save_encodings(encodings, filenames, args.out, args.log)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue