Spaces:
Runtime error
Runtime error
| from image_helpers import convert_images_to_grayscale, crop_center_largest_contour, fetch_similar | |
| import datasets as ds | |
| import re | |
| import torchvision.transforms as T | |
| from transformers import AutoModel, AutoFeatureExtractor | |
| import torch | |
| import random | |
| def similarity_inference(directory): | |
| convert_images_to_grayscale(directory) | |
| crop_center_largest_contour(directory) | |
| # define processing variables needed for embedding calculation | |
| root_directory = "data/" #"C:/Users/josie/OneDrive - Chalmers/Documents/Speckle hackathon/data/" | |
| model_ckpt = "nateraw/vit-base-beans" ## FIND DIFFERENT MODEL | |
| candidate_subset_emb = ds.load_dataset("canadianjosieharrison/2024hackathonembeddingdb")['train'] | |
| extractor = AutoFeatureExtractor.from_pretrained(model_ckpt) | |
| model = AutoModel.from_pretrained(model_ckpt) | |
| transformation_chain = T.Compose( | |
| [ | |
| # We first resize the input image to 256x256 and then we take center crop. | |
| T.Resize(int((256 / 224) * extractor.size["height"])), | |
| T.CenterCrop(extractor.size["height"]), | |
| T.ToTensor(), | |
| T.Normalize(mean=extractor.image_mean, std=extractor.image_std), | |
| ]) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| pt_directory = root_directory + "embedding_db.pt" #"materials/embedding_db.pt" | |
| all_candidate_embeddings = torch.load(pt_directory, map_location=device, weights_only=True) | |
| candidate_ids = [] | |
| for id in range(len(candidate_subset_emb)): | |
| # Create a unique indentifier. | |
| entry = str(id) + "_" + str(random.random()).split('.')[1] | |
| candidate_ids.append(entry) | |
| # load all components | |
| test_ds = ds.load_dataset("imagefolder", data_dir=directory) | |
| label_filenames = ds.load_dataset("imagefolder", data_dir=directory).cast_column("image", ds.Image(decode=False)) | |
| # loop through each component and return top 3 most similar | |
| match_dict = {"ceiling": [], "floor": [], "wall": []} | |
| for i, each_component in enumerate(test_ds['train']): | |
| query_image = each_component["image"] | |
| component_label = label_filenames['train'][i]['image']['path'].split('_')[-1] | |
| print(component_label) | |
| match = re.search(r"([a-zA-Z]+)\d*\.png", component_label) | |
| component_label = match.group(1) | |
| sim_ids = fetch_similar(query_image, transformation_chain, device, model, all_candidate_embeddings, candidate_ids) | |
| for each_match in sim_ids: | |
| texture_filename = candidate_subset_emb[each_match]['filenames'] | |
| image_url = f'https://cdn.polyhaven.com/asset_img/thumbs/{texture_filename}?width=256&height=256' | |
| match_dict[component_label].append(image_url) | |
| return match_dict |