|
import os |
|
import pickle |
|
from .feature_extractor import FeatureExtractor |
|
import time |
|
from tqdm import tqdm |
|
|
|
def precompute_embeddings(): |
|
|
|
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
|
image_dir = os.path.join(base_dir, 'data', 'images') |
|
output_path = os.path.join(base_dir, 'data', 'embeddings.pkl') |
|
|
|
|
|
os.makedirs(image_dir, exist_ok=True) |
|
os.makedirs(os.path.dirname(output_path), exist_ok=True) |
|
|
|
|
|
extractor = FeatureExtractor() |
|
embeddings = [] |
|
image_paths = [] |
|
|
|
valid_images = [f for f in os.listdir(image_dir) |
|
if f.lower().endswith(('.png', '.jpg', '.jpeg'))] |
|
total_images = len(valid_images) |
|
|
|
print(f"\nFound {total_images} images to process") |
|
|
|
start_time = time.time() |
|
for idx, filename in enumerate(tqdm(valid_images, desc="Processing images")): |
|
img_path = os.path.join(image_dir, filename) |
|
try: |
|
embedding = extractor.extract_features(img_path) |
|
embeddings.append(embedding) |
|
image_paths.append(img_path) |
|
except Exception as e: |
|
print(f"\nError processing {filename}: {e}") |
|
|
|
with open(output_path, 'wb') as f: |
|
pickle.dump({'embeddings': embeddings, 'image_paths': image_paths}, f) |
|
|
|
print(f"\nProcessing complete!") |
|
print(f"Successfully processed {len(embeddings)}/{total_images} images") |
|
|
|
return embeddings, image_paths |