Spaces:
Running
Running
File size: 5,805 Bytes
0f10828 |
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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
import numpy as np
import cv2
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.preprocessing import image
from skimage.metrics import structural_similarity as ssim
import os
import argparse
class ImageCharacterClassifier:
def __init__(self, similarity_threshold=0.7):
# Initialize ResNet50 model without top classification layer
self.model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
self.similarity_threshold = similarity_threshold
def load_and_preprocess_image(self, image_path, target_size=(224, 224)):
# Load and preprocess image for ResNet50
img = image.load_img(image_path, target_size=target_size)
img_array = image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
img_array = preprocess_input(img_array)
return img_array
def extract_features(self, image_path):
# Extract deep features using ResNet50
preprocessed_img = self.load_and_preprocess_image(image_path)
features = self.model.predict(preprocessed_img)
return features
def calculate_ssim(self, img1_path, img2_path):
# Calculate SSIM between two images
img1 = cv2.imread(img1_path)
img2 = cv2.imread(img2_path)
# Convert to grayscale if images are in color
if len(img1.shape) == 3:
img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
if len(img2.shape) == 3:
img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
# Resize images to same dimensions
img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))
score = ssim(img1, img2)
return score
def classify_images(self, reference_image_path, image_folder_path):
# Extract features from reference image
reference_features = self.extract_features(reference_image_path)
results = []
# Process each image in the folder
for image_name in os.listdir(image_folder_path):
if image_name.lower().endswith(('.png', '.jpg', '.jpeg')):
image_path = os.path.join(image_folder_path, image_name)
try:
# Calculate SSIM
ssim_score = self.calculate_ssim(reference_image_path, image_path)
# Extract features and calculate similarity
image_features = self.extract_features(image_path)
# Calculate cosine similarity
feature_similarity = np.dot(reference_features.flatten(),
image_features.flatten()) / (
np.linalg.norm(reference_features) * np.linalg.norm(image_features))
# Give more weight to feature similarity
combined_similarity = (0.3 * ssim_score + 0.7 * feature_similarity)
# Classify based on similarity threshold
is_similar = combined_similarity >= self.similarity_threshold
results.append({
'image_name': image_name,
'ssim_score': ssim_score,
'feature_similarity': feature_similarity,
'combined_similarity': combined_similarity,
'is_similar': is_similar
})
except Exception as e:
print(f"Error processing {image_name}: {str(e)}")
continue
return results
def main():
# Create argument parser
parser = argparse.ArgumentParser(description='Image Character Classification')
parser.add_argument('--reference', '-r',
type=str,
required=True,
help='Path to reference image')
parser.add_argument('--folder', '-f',
type=str,
required=True,
help='Path to folder containing images to compare')
parser.add_argument('--threshold', '-t',
type=float,
default=0.5, # Lowered the default threshold
help='Similarity threshold (default: 0.5)')
# Parse arguments
args = parser.parse_args()
# Initialize classifier
classifier = ImageCharacterClassifier(similarity_threshold=args.threshold)
# Check if paths exist
if not os.path.exists(args.reference):
print(f"Error: Reference image not found at {args.reference}")
return
if not os.path.exists(args.folder):
print(f"Error: Image folder not found at {args.folder}")
return
# Perform classification
results = classifier.classify_images(args.reference, args.folder)
# Sort results by similarity score
results.sort(key=lambda x: x['combined_similarity'], reverse=True)
# Print results
print("\nResults sorted by similarity (highest to lowest):")
print("-" * 50)
for result in results:
print(f"\nImage: {result['image_name']}")
print(f"SSIM Score: {result['ssim_score']:.3f}")
print(f"Feature Similarity: {result['feature_similarity']:.3f}")
print(f"Combined Similarity: {result['combined_similarity']:.3f}")
print(f"Is Similar: {result['is_similar']}")
print("-" * 30)
if __name__ == "__main__":
main() |