Spaces:
Sleeping
Sleeping
File size: 2,958 Bytes
73ce1d5 |
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 77 78 79 80 81 82 83 84 85 86 |
import faulthandler
import os
import torch
from tqdm import tqdm
import argparse
import faiss
import PIL
from src.modules.feature_extractor import FeatureExtractor
from src.config.settings import FEATURE_EXTRACTOR_MODELS
faulthandler.enable()
def is_image_file(filename):
valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff', '.webp')
return filename.lower().endswith(valid_extensions)
def main(args=None):
# Set paths
images_dir = "./data/images"
model_dir = "./model"
onnx_path = os.path.join(model_dir, f"{args.feat_extractor}_feature_extractor.onnx")
# Create model directory if it doesn't exist
os.makedirs(model_dir, exist_ok=True)
# Initialize the feature extractor with ONNX support
print(f"Initializing feature extractor with model: {args.feat_extractor}")
feature_extractor = FeatureExtractor(
base_model=args.feat_extractor,
onnx_path=onnx_path
)
# Initialize the vector database indexing
print("Initializing FAISS index...")
index = faiss.IndexFlatIP(feature_extractor.feat_dims)
# Get the list of images in sorted order and filter out non-image files
image_list = sorted([f for f in os.listdir(images_dir) if is_image_file(f)])
print(f"Found {len(image_list)} images to process")
# Limit to first 10 images for initial testing
print(f"Testing with first 10 images")
# Process images and build index
print("Processing images and building index...")
with torch.no_grad():
for img_filename in tqdm(image_list, desc="Processing images"):
try:
# Load image
img = PIL.Image.open(os.path.join(images_dir, img_filename)).convert("RGB")
# Extract features using ONNX model
output = feature_extractor.extract_features(img)
# Prepare features for indexing
output = output.view(output.size(0), -1)
output = output / output.norm(p=2, dim=1, keepdim=True)
output = output.cpu()
# Add to the index
index.add(output.numpy())
except Exception as e:
print(f"Error processing image {img_filename}: {str(e)}")
continue
# Save the index
index_filepath = os.path.join(model_dir, f"db_{args.feat_extractor}.index")
print(f"Saving index to {index_filepath}")
faiss.write_index(index, index_filepath)
print("Index saved successfully!")
if __name__ == "__main__":
# Parse arguments
parser = argparse.ArgumentParser()
parser.add_argument(
"--feat_extractor",
type=str,
default="vit_b_16",
choices=FEATURE_EXTRACTOR_MODELS,
help="Feature extractor model to use"
)
args = parser.parse_args()
# Run the main function
main(args) |