Updated the app to unsqueeze the images

app.py

@@ -8,22 +8,29 @@ from datetime import datetime
 import torch.nn.functional as F
 from typing import List
 
-# Load secrets
-openai_api_key = st.secrets.get("OPENAI_API_KEY")
-# You can now use openai_api_key for anything requiring OpenAI access
-
 # Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load CLIP model + preprocess from OpenAI CLIP
-model, preprocess = clip.load("ViT-L/14", device=device)
+# Load CLIP model and preprocessor (ViT-B/32 = small model, CPU-friendly)
+model, preprocess = clip.load("ViT-B/32", device=device)
 model.eval()
 
-# Ensure reproducibility
-torch.set_grad_enabled(False)
+# Display app title and information
+st.set_page_config(page_title="Few-Shot Fault Detection", layout="wide")
+st.title("🛠️ Few-Shot Fault Detection (Industrial Quality Control)")
+
+st.markdown("""
+This demo uses the **smaller `ViT-B/32` encoder** from OpenAI's CLIP model to classify test images as **Nominal** or **Defective**, based on few-shot learning using user-provided reference images.
+
+⚠️ **Note**: This app is running on a **free CPU tier** and is meant for demonstration purposes. For more advanced use cases, including GPU acceleration, custom training, and larger models, please refer to:
+
+📄 [Megahed et al. (2025)](https://arxiv.org/abs/2501.12596):  
+*Adapting OpenAI's CLIP Model for Few-Shot Image Inspection in Manufacturing Quality Control: An Expository Case Study with Multiple Application Examples*
+
+🔗 [GitHub & Colab links available in the paper](https://arxiv.org/abs/2501.12596)
+""")
 
-# Import the few-shot classification function
-# --- COPY YOUR FUNCTION DEFINITION BELOW DIRECTLY OR PUT IT IN A SEPARATE FILE ---
+# --- Few-shot classification logic ---
 def few_shot_fault_classification(
     test_images: List[Image.Image],
     test_image_filenames: List[str],
@@ -47,16 +54,16 @@ def few_shot_fault_classification(
     results = []
 
     with torch.no_grad():
-        nominal_features = torch.stack([model.encode_image(img).to(device) for img in nominal_images])
+        nominal_features = torch.stack([model.encode_image(img.unsqueeze(0)).squeeze(0).to(device) for img in nominal_images])
         nominal_features /= nominal_features.norm(dim=-1, keepdim=True)
 
-        defective_features = torch.stack([model.encode_image(img).to(device) for img in defective_images])
+        defective_features = torch.stack([model.encode_image(img.unsqueeze(0)).squeeze(0).to(device) for img in defective_images])
         defective_features /= defective_features.norm(dim=-1, keepdim=True)
 
         csv_data = []
 
         for idx, test_img in enumerate(test_images):
-            test_features = model.encode_image(test_img).to(device)
+            test_features = model.encode_image(test_img.unsqueeze(0)).squeeze(0).to(device)
             test_features /= test_features.norm(dim=-1, keepdim=True)
 
             max_nom_sim, max_def_sim = -float('inf'), -float('inf')
@@ -110,7 +117,7 @@ def few_shot_fault_classification(
 
     return ""
 
-# Initialize app state
+# --- App state ---
 if 'nominal_images' not in st.session_state:
     st.session_state.nominal_images = []
 if 'defective_images' not in st.session_state:
@@ -120,16 +127,12 @@ if 'test_images' not in st.session_state:
 if 'results' not in st.session_state:
     st.session_state.results = []
 
-st.set_page_config(page_title="Few-Shot Fault Detection", layout="wide")
-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.")
-
+# --- Tabs ---
 tab1, tab2, tab3 = st.tabs(["📥 Upload Reference Images", "🔍 Test Classification", "📊 Results"])
 
-# --- Tab 1: Upload Reference Images ---
+# Tab 1: Upload Reference Images
 with tab1:
     st.header("Upload Reference Images")
-
     nominal_files = st.file_uploader("Upload Nominal Images", accept_multiple_files=True, type=['png', 'jpg', 'jpeg'])
     defective_files = st.file_uploader("Upload Defective Images", accept_multiple_files=True, type=['png', 'jpg', 'jpeg'])
 
@@ -143,10 +146,9 @@ with tab1:
         st.session_state.defective_descriptions = [file.name for file in defective_files]
         st.success(f"Uploaded {len(defective_files)} defective images.")
 
-# --- Tab 2: Classify Test Images ---
+# Tab 2: Test Classification
 with tab2:
     st.header("Upload Test Image(s)")
-
     test_files = st.file_uploader("Upload Test Images", accept_multiple_files=True, type=['png', 'jpg', 'jpeg'])
 
     if st.button("🔍 Run Classification") and test_files:
@@ -169,10 +171,9 @@ with tab2:
         st.success("Classification complete!")
         st.session_state.results = "streamlit_results.csv"
 
-# --- Tab 3: View/Download Results ---
+# Tab 3: View/Download Results
 with tab3:
     st.header("Classification Results")
-
     if os.path.exists("streamlit_results.csv"):
         df = pd.read_csv("streamlit_results.csv")
         st.dataframe(df)