Upload image_classifier.py

image_classifier.py

@@ -0,0 +1,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()