File size: 1,557 Bytes
3cfebcb |
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 |
import os
import pickle
from .feature_extractor import FeatureExtractor
import time
from tqdm import tqdm
def precompute_embeddings():
# Use absolute paths for Hugging Face Spaces
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')
# Create directories if they don't exist
os.makedirs(image_dir, exist_ok=True)
os.makedirs(os.path.dirname(output_path), exist_ok=True)
# Rest of your existing code...
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 |