Update app.py

app.py

@@ -2,48 +2,256 @@ import streamlit as st
 import torch
 import open_clip
 from PIL import Image
+import os
 from classifier import few_shot_fault_classification
 
-# Load lightweight CLIP model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model, _, preprocess = open_clip.create_model_and_transforms('RN50', pretrained='openai')
-model = model.to(device)
-model.eval()
+# Set page configuration
+st.set_page_config(
+    page_title="Industrial QC Image Classifier",
+    page_icon="🔍",
+    layout="wide"
+)
 
-st.title("🛠️ Few-Shot Fault Detection (Industrial Quality Control)")
+# Initialize session state variables if they don't exist
+if 'nominal_images' not in st.session_state:
+    st.session_state.nominal_images = []
+    st.session_state.nominal_descriptions = []
+    st.session_state.nominal_filenames = []
+    
+if 'defective_images' not in st.session_state:
+    st.session_state.defective_images = []
+    st.session_state.defective_descriptions = []
+    st.session_state.defective_filenames = []
+    
+if 'test_results' not in st.session_state:
+    st.session_state.test_results = []
 
-st.markdown("Upload **10 Nominal Images**, **10 Defective Images**, and one or more **Test Images** to classify.")
+if 'model' not in st.session_state:
+    st.session_state.device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Load model
+    with st.spinner("Loading CLIP model..."):
+        model, _, preprocess = open_clip.create_model_and_transforms('RN50', pretrained='openai')
+        model = model.to(st.session_state.device)
+        model.eval()
+        st.session_state.model = model
+        st.session_state.preprocess = preprocess
 
-col1, col2 = st.columns(2)
-with col1:
-    nominal_files = st.file_uploader("Upload Nominal Images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
-with col2:
-    defective_files = st.file_uploader("Upload Defective Images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
+def process_uploaded_image(uploaded_file):
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        return image
+    return None
 
-test_files = st.file_uploader("Upload Test Images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
+def display_image_grid(images, captions, columns=5):
+    if not images:
+        return
+    
+    rows = (len(images) + columns - 1) // columns
+    
+    for i in range(rows):
+        cols = st.columns(columns)
+        for j in range(columns):
+            idx = i * columns + j
+            if idx < len(images):
+                with cols[j]:
+                    st.image(images[idx], caption=captions[idx], use_column_width=True)
 
-if st.button("Classify Test Images"):
-    if len(nominal_files) < 1 or len(defective_files) < 1 or len(test_files) < 1:
-        st.warning("Please upload at least 1 image in each category.")
-    else:
-        st.info("Running classification...")
+def main():
+    st.title("🛠️ Few-Shot Fault Detection (Industrial Quality Control)")
+    st.markdown("Upload **Nominal Images** (good parts), **Defective Images** (bad parts), and **Test Images** to classify.")
+    
+    # Create tabs for different functionality
+    tab1, tab2, tab3 = st.tabs(["📥 Upload Reference Images", "🔍 Test Classification", "📊 Results"])
+    
+    with tab1:
+        st.header("Upload Reference Images")
+        
+        # Create two columns for nominal and defective images
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            st.subheader("Good Parts (Nominal)")
+            
+            # Input for nominal description
+            nominal_desc = st.text_input("Nominal Description", 
+                                       value="Good part without defects")
+            
+            # File uploader for nominal images
+            nominal_files = st.file_uploader("Upload images of good parts (10 recommended)", 
+                                          type=["jpg", "jpeg", "png"], 
+                                          accept_multiple_files=True,
+                                          key="nominal_upload")
+            
+            if nominal_files and st.button("Add Nominal Images"):
+                for file in nominal_files:
+                    img = process_uploaded_image(file)
+                    if img:
+                        # Preprocess the image for the model
+                        preprocessed_img = st.session_state.preprocess(img).unsqueeze(0)
+                        st.session_state.nominal_images.append(preprocessed_img)
+                        st.session_state.nominal_descriptions.append(nominal_desc)
+                        st.session_state.nominal_filenames.append(file.name)
+                
+                st.success(f"Added {len(nominal_files)} nominal images!")
+                
+            # Display current nominal image count
+            st.write(f"Current nominal images: {len(st.session_state.nominal_images)}")
+            
+            # Display nominal images in a grid if we have any
+            if st.session_state.nominal_images and st.session_state.nominal_filenames:
+                st.subheader("Current Nominal Images")
+                # We'll display the filenames instead of the actual tensor images
+                display_placeholders = [f"Image {i+1}" for i in range(len(st.session_state.nominal_filenames))]
+                st.write(", ".join(st.session_state.nominal_filenames))
+                
+                if st.button("Clear Nominal Images"):
+                    st.session_state.nominal_images = []
+                    st.session_state.nominal_descriptions = []
+                    st.session_state.nominal_filenames = []
+                    st.success("Cleared all nominal images!")
+            
+        with col2:
+            st.subheader("Defective Parts")
+            
+            # Input for defective description
+            defective_desc = st.text_input("Defective Description", 
+                                        value="Part with visible defects")
+            
+            # File uploader for defective images
+            defective_files = st.file_uploader("Upload images of defective parts (10 recommended)", 
+                                           type=["jpg", "jpeg", "png"], 
+                                           accept_multiple_files=True,
+                                           key="defective_upload")
+            
+            if defective_files and st.button("Add Defective Images"):
+                for file in defective_files:
+                    img = process_uploaded_image(file)
+                    if img:
+                        # Preprocess the image for the model
+                        preprocessed_img = st.session_state.preprocess(img).unsqueeze(0)
+                        st.session_state.defective_images.append(preprocessed_img)
+                        st.session_state.defective_descriptions.append(defective_desc)
+                        st.session_state.defective_filenames.append(file.name)
+                
+                st.success(f"Added {len(defective_files)} defective images!")
+                
+            # Display current defective image count
+            st.write(f"Current defective images: {len(st.session_state.defective_images)}")
+            
+            # Display defective images in a grid if we have any
+            if st.session_state.defective_images and st.session_state.defective_filenames:
+                st.subheader("Current Defective Images")
+                # We'll display the filenames instead of the actual tensor images
+                st.write(", ".join(st.session_state.defective_filenames))
+                
+                if st.button("Clear Defective Images"):
+                    st.session_state.defective_images = []
+                    st.session_state.defective_descriptions = []
+                    st.session_state.defective_filenames = []
+                    st.success("Cleared all defective images!")
+    
+    with tab2:
+        st.header("Test Image Classification")
+        
+        # Check if we have enough reference images
+        if len(st.session_state.nominal_images) == 0 or len(st.session_state.defective_images) == 0:
+            st.warning("You need to upload at least one nominal image and one defective image before testing.")
+        else:
+            st.write("Upload a test image to classify it as nominal or defective.")
+            
+            # File uploader for test image
+            test_files = st.file_uploader("Upload test images", 
+                                       type=["jpg", "jpeg", "png"],
+                                       accept_multiple_files=True,
+                                       key="test_upload")
+            
+            if test_files:
+                test_images = []
+                test_image_displays = []
+                test_filenames = []
+                
+                for file in test_files:
+                    img = process_uploaded_image(file)
+                    if img:
+                        test_image_displays.append(img)
+                        # Preprocess the image for the model
+                        preprocessed_img = st.session_state.preprocess(img).unsqueeze(0)
+                        test_images.append(preprocessed_img.squeeze(0))
+                        test_filenames.append(file.name)
+                
+                # Display the test images
+                if test_image_displays:
+                    display_image_grid(
+                        test_image_displays, 
+                        [f"Test: {name}" for name in test_filenames]
+                    )
+                    
+                    if st.button("Classify Images"):
+                        with st.spinner("Classifying images..."):
+                            # Create output directory
+                            os.makedirs("./results", exist_ok=True)
+                            
+                            # Run classification
+                            results = few_shot_fault_classification(
+                                model=st.session_state.model,
+                                test_images=test_images,
+                                test_image_filenames=test_filenames,
+                                nominal_images=[img.squeeze(0) for img in st.session_state.nominal_images],
+                                nominal_descriptions=st.session_state.nominal_descriptions,
+                                defective_images=[img.squeeze(0) for img in st.session_state.defective_images],
+                                defective_descriptions=st.session_state.defective_descriptions,
+                                num_few_shot_nominal_imgs=len(st.session_state.nominal_images),
+                                device=st.session_state.device,
+                                file_path="./results",
+                                file_name="classification_results.csv"
+                            )
+                            
+                            for i, result in enumerate(results):
+                                st.session_state.test_results.append({
+                                    "image": test_image_displays[i],
+                                    "filename": test_filenames[i],
+                                    **result
+                                })
+                            
+                            # Display classification results
+                            for result in results:
+                                classification = result["classification_result"]
+                                color = "green" if classification == "Nominal" else "red"
+                                
+                                st.markdown(f"### {result['image_name']}: <span style='color:{color}'>{classification}</span>", unsafe_allow_html=True)
+                                
+                                # Display probabilities
+                                col1, col2 = st.columns(2)
+                                with col1:
+                                    st.metric("Good Part Probability", f"{result['non_defect_prob']:.3f}")
+                                with col2:
+                                    st.metric("Defective Part Probability", f"{result['defect_prob']:.3f}")
+    
+    with tab3:
+        st.header("Classification Results")
+        
+        if not st.session_state.test_results:
+            st.info("No classification results yet. Test some images in the 'Test Classification' tab.")
+        else:
+            st.write(f"Total images classified: {len(st.session_state.test_results)}")
+            
+            # Display results in a table
+            results_df = [{
+                "Filename": r["filename"],
+                "Classification": r["classification_result"],
+                "Good Prob": f"{r['non_defect_prob']:.3f}",
+                "Defect Prob": f"{r['defect_prob']:.3f}",
+                "Time": r["datetime_of_operation"].split('T')[0]
+            } for r in st.session_state.test_results]
+            
+            st.table(results_df)
+            
+            # Option to clear results
+            if st.button("Clear Results"):
+                st.session_state.test_results = []
+                st.success("Classification results cleared!")
 
-        nominal_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0).to(device) for f in nominal_files]
-        defective_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0).to(device) for f in defective_files]
-        test_imgs = [preprocess(Image.open(f).convert("RGB")).unsqueeze(0).to(device) for f in test_files]
-
-        results = few_shot_fault_classification(
-            model=model,
-            test_images=[img.squeeze(0) for img in test_imgs],
-            test_image_filenames=[f.name for f in test_files],
-            nominal_images=[img.squeeze(0) for img in nominal_imgs],
-            nominal_descriptions=[f.name for f in nominal_files],
-            defective_images=[img.squeeze(0) for img in defective_imgs],
-            defective_descriptions=[f.name for f in defective_files],
-            num_few_shot_nominal_imgs=len(nominal_files),
-            device=device
-        )
-
-        for res in results:
-            st.write(f"**{res['image_path']}** ➜ {res['classification_result']} "
-                     f"(Nominal: {res['non_defect_prob']}, Defective: {res['defect_prob']})")
+if __name__ == "__main__":
+    main()