Create app.py

app.py

@@ -0,0 +1,339 @@
+import streamlit as st
+import cv2
+from PIL import Image, ImageEnhance
+import numpy as np
+import time
+from skimage.metrics import structural_similarity as ssim
+import base64
+from datetime import datetime
+import torch
+
+# Load pre-trained YOLOv5 model for object detection
+@st.cache_resource
+def load_yolo_model():
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
+    return model
+
+def load_css():
+    st.markdown("""
+        <style>
+        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&display=swap');
+        
+        .stApp {
+            background: linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%);
+            font-family: 'Inter', sans-serif;
+            color: #e0e0e0;
+        }
+        
+        .main {
+            padding: 2rem;
+            max-width: 1200px;
+            margin: 0 auto;
+        }
+        
+        .stButton>button {
+            background: linear-gradient(135deg, #2196F3 0%, #1976D2 100%);
+            color: white;
+            padding: 0.75rem 1.5rem;
+            border-radius: 10px;
+            border: none;
+            box-shadow: 0 4px 6px rgba(0,0,0,0.2);
+            transition: all 0.3s ease;
+            font-weight: 500;
+            letter-spacing: 0.5px;
+        }
+        
+        .stButton>button:hover {
+            transform: translateY(-2px);
+            box-shadow: 0 6px 12px rgba(0,0,0,0.3);
+        }
+        
+        .upload-container {
+            background: #2d2d2d;
+            border-radius: 15px;
+            padding: 1.5rem;
+            box-shadow: 0 4px 6px rgba(0,0,0,0.2);
+            transition: all 0.3s ease;
+            margin-bottom: 1rem;
+        }
+        
+        .upload-container:hover {
+            box-shadow: 0 6px 12px rgba(0,0,0,0.3);
+        }
+        
+        .upload-box {
+            border: 2px dashed #404040;
+            border-radius: 12px;
+            padding: 2rem;
+            text-align: center;
+            background: #333333;
+            transition: all 0.3s ease;
+            cursor: pointer;
+        }
+        
+        .upload-box:hover {
+            border-color: #2196F3;
+            background: #383838;
+        }
+        
+        .results-container {
+            background: #2d2d2d;
+            border-radius: 15px;
+            padding: 2rem;
+            box-shadow: 0 4px 6px rgba(0,0,0,0.2);
+            color: #e0e0e0;
+        }
+        
+        .metric-card {
+            background: #333333;
+            border-radius: 10px;
+            padding: 1rem;
+            margin: 0.5rem 0;
+            border-left: 4px solid #2196F3;
+            color: #e0e0e0;
+        }
+        
+        .stProgress > div > div {
+            background: linear-gradient(90deg, #2196F3, #64B5F6);
+            border-radius: 10px;
+        }
+        
+        @keyframes pulse {
+            0% { opacity: 1; }
+            50% { opacity: 0.5; }
+            100% { opacity: 1; }
+        }
+        
+        .loading {
+            animation: pulse 1.5s infinite;
+        }
+        </style>
+    """, unsafe_allow_html=True)
+
+def enhance_image(image):
+    """
+    Basic image enhancement with default settings
+    """
+    enhancer = ImageEnhance.Brightness(image)
+    image = enhancer.enhance(1.0)
+    enhancer = ImageEnhance.Contrast(image)
+    image = enhancer.enhance(1.0)
+    enhancer = ImageEnhance.Sharpness(image)
+    image = enhancer.enhance(1.0)
+    return image
+
+def compare_images(img1, img2, progress_bar):
+    """
+    Compare two images and return the processed image, similarity score, and difference percentage
+    """
+    try:
+        progress_bar.progress(0)
+        
+        # Convert images to numpy arrays and ensure same size
+        img1 = np.array(img1.resize(img2.size))
+        img2 = np.array(img2)
+        progress_bar.progress(20)
+        
+        # Normalize images
+        img1 = cv2.normalize(img1, None, 0, 255, cv2.NORM_MINMAX)
+        img2 = cv2.normalize(img2, None, 0, 255, cv2.NORM_MINMAX)
+        
+        # Convert to grayscale
+        gray1 = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
+        gray2 = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)
+        progress_bar.progress(40)
+        
+        # Calculate SSIM
+        score, diff = ssim(gray1, gray2, full=True)
+        progress_bar.progress(60)
+        
+        # Generate heatmap
+        diff = (diff * 255).astype(np.uint8)
+        heatmap = cv2.applyColorMap(diff, cv2.COLORMAP_JET)
+        progress_bar.progress(80)
+        
+        # Highlight differences in red color
+        diff_mask = cv2.absdiff(gray1, gray2)
+        diff_mask = cv2.cvtColor(diff_mask, cv2.COLOR_GRAY2RGB)
+        diff_mask[np.where((diff_mask == [255, 255, 255]).all(axis=2))] = [0, 0, 255]  # Red color for differences
+        
+        # Combine original image with difference mask
+        result_img = cv2.addWeighted(img1, 0.7, diff_mask, 0.3, 0)
+        
+        # Calculate pixel-wise differences
+        diff_percentage = (np.count_nonzero(diff_mask[:, :, 2] > 0) / (diff_mask.shape[0] * diff_mask.shape[1])) * 100
+        
+        # Ensure that the difference percentage is consistent with the similarity score
+        diff_percentage = 100 - (score * 100)
+        
+        progress_bar.progress(100)
+        
+        return result_img, score, diff_percentage, heatmap
+        
+    except Exception as e:
+        st.error(f"Error comparing images: {str(e)}")
+        return None, 0, 0, None
+
+def detect_objects(image, model):
+    """
+    Perform object detection on the image using YOLOv5
+    """
+    try:
+        results = model(image)
+        results_df = results.pandas().xyxy[0]
+        return results_df
+    except Exception as e:
+        st.error(f"Error in object detection: {str(e)}")
+        return None
+
+def draw_object_boxes(image, objects_df):
+    """
+    Draw bounding boxes on the image for detected objects
+    """
+    for _, row in objects_df.iterrows():
+        xmin, ymin, xmax, ymax, confidence, class_name = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax']), row['confidence'], row['name']
+        # Draw bounding box
+        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
+        # Add label
+        cv2.putText(image, f"{class_name} {confidence:.2f}", (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+    return image
+
+def main():
+    load_css()
+    
+    # Initialize session state for results
+    if "results" not in st.session_state:
+        st.session_state.results = None
+    
+    # Load YOLOv5 model
+    yolo_model = load_yolo_model()
+    
+    # App header
+    st.markdown("""
+        <div style='text-align: center; margin-bottom: 2rem; background: linear-gradient(135deg, #2196F3 0%, #1976D2 100%); padding: 2rem; border-radius: 15px; color: white;'>
+            <h1 style='margin: 0;'>🔍 Image Comparison Tool</h1>
+            <p style='margin: 1rem 0 0 0; opacity: 0.9;'>Compare images, highlight differences, and detect objects</p>
+        </div>
+    """, unsafe_allow_html=True)
+    
+    # Main content for image upload and display
+    st.markdown("<div class='upload-container'>", unsafe_allow_html=True)
+    st.markdown("### 📁 Upload Images")
+    
+    col1, col2 = st.columns(2)
+    
+    # Reference Image Upload
+    with col1:
+        reference_image = st.file_uploader(
+            "Drop or select reference image",
+            type=["jpg", "jpeg", "png"],
+            key="reference"
+        )
+        if reference_image:
+            img1 = Image.open(reference_image)
+            img1 = enhance_image(img1)
+            st.image(img1, caption="Reference Image", use_column_width=True)
+            # Clear previous results when a new image is uploaded
+            st.session_state.results = None
+    
+    # New Image Upload
+    with col2:
+        new_image = st.file_uploader(
+            "Drop or select comparison image",
+            type=["jpg", "jpeg", "png"],
+            key="new"
+        )
+        if new_image:
+            img2 = Image.open(new_image)
+            img2 = enhance_image(img2)
+            st.image(img2, caption="Comparison Image", use_column_width=True)
+            # Clear previous results when a new image is uploaded
+            st.session_state.results = None
+    
+    st.markdown("</div>", unsafe_allow_html=True)
+    
+    # Sidebar for results and download
+    st.sidebar.markdown("### 🎯 Analysis Results")
+    
+    if reference_image and new_image:
+        compare_button = st.sidebar.button("🔍 Analyze Images", use_container_width=True)
+        
+        if compare_button or st.session_state.results:
+            if not st.session_state.results:
+                with st.spinner("Processing images..."):
+                    progress_bar = st.sidebar.progress(0)
+                    
+                    start_time = time.time()
+                    result_img, score, diff_percentage, heatmap = compare_images(img1, img2, progress_bar)
+                    processing_time = time.time() - start_time
+                    
+                    # Perform object detection
+                    objects_df = detect_objects(result_img, yolo_model)
+                    
+                    # Draw bounding boxes on the analyzed image
+                    if objects_df is not None:
+                        result_img = draw_object_boxes(result_img, objects_df)
+                    
+                    # Store results in session state
+                    st.session_state.results = {
+                        "result_img": result_img,
+                        "heatmap": heatmap,
+                        "score": score,
+                        "diff_percentage": diff_percentage,
+                        "processing_time": processing_time,
+                        "objects_df": objects_df
+                    }
+            
+            # Display analyzed image (processed image with differences highlighted) in sidebar
+            st.sidebar.image(st.session_state.results["result_img"], caption="Analyzed Image (Differences Highlighted)", use_column_width=True)
+            
+            # Display heatmap in sidebar
+            st.sidebar.image(st.session_state.results["heatmap"], caption="Heatmap", use_column_width=True)
+            
+            # Display metrics in sidebar
+            st.sidebar.markdown("### 📊 Metrics")
+            st.sidebar.markdown(f"""
+                <div class='metric-card'>
+                    <h4>Similarity Score</h4>
+                    <h2 style='color: #2196F3'>{st.session_state.results["score"]:.2%}</h2>
+                </div>
+            """, unsafe_allow_html=True)
+            
+            st.sidebar.markdown(f"""
+                <div class='metric-card'>
+                    <h4>Difference Detected</h4>
+                    <h2 style='color: #2196F3'>{st.session_state.results["diff_percentage"]:.2f}%</h2>
+                </div>
+            """, unsafe_allow_html=True)
+            
+            st.sidebar.markdown(f"""
+                <div class='metric-card'>
+                    <h4>Processing Time</h4>
+                    <h2 style='color: #2196F3'>{st.session_state.results["processing_time"]:.2f}s</h2>
+                </div>
+            """, unsafe_allow_html=True)
+            
+            # Display detected objects
+            if st.session_state.results["objects_df"] is not None:
+                st.sidebar.markdown("### 🔍 Detected Objects")
+                st.sidebar.dataframe(st.session_state.results["objects_df"])
+            
+            # Download analyzed image
+            st.sidebar.markdown("### 📥 Download Analyzed Image")
+            st.sidebar.download_button(
+                "Download Analyzed Image",
+                data=cv2.imencode('.png', cv2.cvtColor(st.session_state.results["result_img"], cv2.COLOR_RGB2BGR))[1].tobytes(),
+                file_name=f"analyzed_image_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png",
+                mime="image/png"
+            )
+    
+    # Footer
+    st.markdown("""
+        <div style='text-align: center; margin-top: 2rem; padding: 1rem; background: #2d2d2d; border-radius: 10px; box-shadow: 0 4px 6px rgba(0,0,0,0.2);'>
+            <p style='color: #888; margin: 0;'>Built with ❤️ using Streamlit | Last updated: December 2024</p>
+            <p style='color: #888; font-size: 0.9em; margin: 0.5rem 0 0 0;'>Image Comparison Tool v1.0</p>
+        </div>
+    """, unsafe_allow_html=True)
+
+if __name__ == "__main__":
+    main()