Last active
October 20, 2023 07:32
-
-
Save amenabe22/a99b1a886cf2b4536620e2552b0ed0fa to your computer and use it in GitHub Desktop.
bg-similarity checker
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import os | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from skimage.metrics import structural_similarity as ssim | |
| # Load the YOLOv5 model | |
| model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s') | |
| # Load your three images | |
| image_path1 = 'image1.jpg' | |
| image_path2 = 'image2.jpg' | |
| image_path3 = 'colors.jpg' | |
| # path where to save the background images for better preview | |
| BG_PREV_PATH = "/Users/amenmohammed/Documents/projects/artisse-ai/altered" | |
| def extract_background(image_path, target_size=(640, 640)): | |
| img = Image.open(image_path) | |
| # Perform object detection with YOLOv5 | |
| # results = detect.detect_image(model, img) | |
| results = model(img) | |
| prediction = results.pred[0] | |
| # resize images to a common size | |
| img = img.resize(target_size, Image.LANCZOS) | |
| # Extract the background by masking the objects | |
| background = img | |
| for det in prediction: | |
| x1, y1, x2, y2, class_id, confidence = det | |
| x1, y1, x2, y2 = map(int, (x1, y1, x2, y2)) | |
| background = np.array(background) | |
| background[y1:y2, x1:x2, :] = 0 # Black out the detected object | |
| return background | |
| if __name__ == "__main__": | |
| # Set a common target size for all images | |
| target_size = (640, 640) | |
| # Extract backgrounds | |
| background1 = extract_background(image_path1, target_size) | |
| background2 = extract_background(image_path2, target_size) | |
| background3 = extract_background(image_path3, target_size) | |
| # Save the extracted backgrounds as files | |
| background1_pil = Image.fromarray(background1) | |
| background2_pil = Image.fromarray(background2) | |
| background3_pil = Image.fromarray(background3) | |
| background1_pil.save(os.path.join(BG_PREV_PATH,"bg1.jpg")) | |
| background2_pil.save(os.path.join(BG_PREV_PATH,"bg2.jpg")) | |
| background3_pil.save(os.path.join(BG_PREV_PATH,"bg3.jpg")) | |
| # Calculate structural similarity (SSIM) between the background images | |
| ssim_score12 = ssim(background1, background2, channel_axis=-1, data_range=255) | |
| ssim_score13 = ssim(background1, background3, channel_axis=-1, data_range=255) | |
| # Calculate the average SSIM as a similarity score | |
| similarity_score = (ssim_score12 + ssim_score13) / 2 * 100 | |
| # Print the result | |
| print(f"Background Similarity: {similarity_score:.2f}%") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment