Spaces:
Sleeping
Sleeping
| import torch | |
| def get_image_features(input_image, feature_extractor): | |
| with feature_extractor.no_bar(), feature_extractor.no_logging(): | |
| test_dl = feature_extractor.dls.test_dl([input_image]) | |
| inp, features, _, dec = feature_extractor.get_preds(dl=test_dl, with_input=True, with_decoded=True) | |
| return features | |
| def get_similar_image(input_image, feature_extractor, features_tensor, image_paths): | |
| # Get the features of the input image | |
| user_features = get_image_features(input_image, feature_extractor) | |
| user_features = user_features.view(1, -1) # Reshape to 2D tensor | |
| # Compute cosine similarity | |
| similarity_scores = torch.nn.functional.cosine_similarity(user_features, features_tensor) | |
| # Get the index of the most similar image | |
| most_similar_index = torch.argmax(similarity_scores) | |
| # Get the path of the most similar image | |
| most_similar_image_path = image_paths[most_similar_index] | |
| # Get the maximum similarity score | |
| max_similarity = torch.max(similarity_scores) | |
| return most_similar_image_path, max_similarity | |
| def plot_side_by_side(input_image, similar_image, show=True, save_path=None): | |
| similar_image_thumb = similar_image.to_thumb(224) | |
| user_image_thumb = input_image.to_thumb(224) | |
| # Create a figure with two subplots | |
| fig, (ax1, ax2) = plt.subplots(1, 2) | |
| # Display the images | |
| ax1.imshow(similar_image_thumb) | |
| ax2.imshow(user_image_thumb) | |
| # Optionally, remove the axes for a cleaner look | |
| ax1.axis('off') | |
| ax2.axis('off') | |
| fig.suptitle('Is It Really Worth It?', fontsize=20, weight='bold') | |
| if save_path: | |
| plt.savefig(save_path) | |
| plt.close() | |
| if show: | |
| plt.show() | |
| def test_model(feature_extractor, features_tensor, model_name, image_paths, input_dir=Path('input'), output_dir=Path('output'), show=False): | |
| save_dir = output_dir / model_name | |
| save_dir.mkdir(parents=True, exist_ok=True) | |
| for input_path in input_dir.iterdir(): | |
| save_path = os.path.join(save_dir, os.path.splitext(os.path.basename(input_path))[0] + '.jpg') | |
| input_image = PILImage.create(input_path) | |
| process_image(input_image, feature_extractor, features_tensor, image_paths, save_path=save_path, show=show) | |
| def random_crop(input_image, scale=(0.3, 0.4)): | |
| width, height = input_image.size | |
| # Calculate random width and height | |
| new_width = random.randint(int(width * scale[0]), int(width * scale[1])) | |
| new_height = random.randint(int(height * scale[0]), int(height * scale[1])) | |
| # Calculate random position for the crop | |
| left = random.randint(0, width - new_width) | |
| top = random.randint(0, height - new_height) | |
| # Perform the crop | |
| cropped_img = input_image.crop((left, top, left + new_width, top + new_height)) | |
| # Resize the cropped image to 224x224 | |
| resized_img = cropped_img.resize((224, 224)) | |
| # Return the resized image and its coordinates | |
| return resized_img, (left, top, left + new_width, top + new_height) | |
| def process_image(input_image, feature_extractor, features_tensor, image_paths, show=True, save_path=None): | |
| max_similarity = -1 | |
| most_similar_image_path = None | |
| input_image_crop_coords = None | |
| reference_image_crop_coords = None | |
| # Apply the transform 10 times to get 10 random crops | |
| for i in range(10): | |
| # Perform a random crop | |
| cropped_img, crop_coords = random_crop(input_image) | |
| # Get the most similar image for the cropped image and its similarity score | |
| similar_image_path, similarity = get_similar_image(cropped_img, feature_extractor, features_tensor, image_paths) | |
| # If this image is more similar than the previous ones, keep it | |
| if similarity > max_similarity: | |
| max_similarity = similarity | |
| most_similar_image_path = similar_image_path | |
| input_image_crop_coords = crop_coords | |
| reference_image_crop_coords = get_crop_coords_from_filename(similar_image_path) | |
| # Get the parent and crop coordinates from the filename | |
| parent, filename = os.path.split(most_similar_image_path) | |
| # Plot the input image and the most similar image side by side | |
| plot_side_by_side(input_image, PILImage.create(most_similar_image_path), input_image_crop_coords, reference_image_crop_coords, show=show, save_path=save_path) | |
| return parent, filename, input_image_crop_coords, reference_image_crop_coords |