Update app.py

app.py

@@ -1,93 +1,95 @@
-import os
-import torch
-import numpy as np
-import gradio as gr
-from PIL import Image
-from transformers import CLIPProcessor, CLIPModel
-
-# Load model parameters
-params = torch.load("clip_classification_params.pth", weights_only=False)
-
-correct_centroid = torch.tensor(params["correct_centroid"])
-incorrect_centroid = torch.tensor(params["incorrect_centroid"])
-threshold = params["threshold"]
-
-# Load CLIP model
-MODEL_NAME = "openai/clip-vit-base-patch32"
-clip_model = CLIPModel.from_pretrained(MODEL_NAME)
-clip_processor = CLIPProcessor.from_pretrained(MODEL_NAME)
-clip_model.eval()
-
-# Paths to example images
-correct_examples = [
-    "dataset/correct/correct (1).webp",
-    "dataset/correct/correct (2).webp",
-    "dataset/correct/correct (3).webp"
-]
-
-incorrect_examples = [
-    "dataset/incorrect/incorrect (1).webp",
-    "dataset/incorrect/incorrect (2).webp",
-    "dataset/incorrect/incorrect (3).webp"
-]
-
-# Function to classify an image
-def classify_installation(image):
-    """Classify if the bed installation is correct or incorrect."""
-    inputs = clip_processor(images=image, return_tensors="pt", padding=True)
-    with torch.no_grad():
-        embedding = clip_model.get_image_features(**inputs)
-        embedding = embedding / embedding.norm(p=2, dim=-1, keepdim=True)
-
-    # Compute similarity to correct centroid
-    similarity = torch.matmul(embedding, correct_centroid)
-    
-    # Compare similarity with threshold
-    if similarity.item() >= threshold:
-        return f"✅ Correct Installation (similarity = {similarity.item():.2f})"
-    else:
-        return f"❌ Incorrect Installation (similarity = {similarity.item():.2f})"
-
-# Function to load image from file path
-def load_image(image_path):
-    return Image.open(image_path).convert("RGB")
-
-# Function to process selected example image
-def process_example(image_path):
-    image = load_image(image_path)
-    return classify_installation(image), image
-
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("# 🛏️ Bed Installation Classifier")
-    gr.Markdown("Upload an image or select one from the examples below to check if the bed installation is correct.")
-
-    with gr.Row():
-        uploaded_image = gr.Image(type="pil", label="Upload an image for testing")
-        output_text = gr.Textbox(label="Result")
-
-    gr.Markdown("### Check Installations (Click Button to Classify)")
-    with gr.Row():
-        correct_buttons = []
-        for i, img_path in enumerate(correct_examples):
-            with gr.Column():
-                gr.Image(value=load_image(img_path), interactive=False, width=150, height=150)
-                btn = gr.Button(value=f"Check Accuracy")
-                correct_buttons.append((btn, img_path))
-
-    with gr.Row():
-        incorrect_buttons = []
-        for i, img_path in enumerate(incorrect_examples):
-            with gr.Column():
-                gr.Image(value=load_image(img_path), interactive=False, width=150, height=150)
-                btn = gr.Button(value=f"Check Accuracy")
-                incorrect_buttons.append((btn, img_path))
-
-    # Connect buttons to classification function
-    for btn, img_path in correct_buttons + incorrect_buttons:
-        btn.click(fn=process_example, inputs=[gr.State(img_path)], outputs=[output_text, uploaded_image])
-
-    # Process uploaded image
-    uploaded_image.change(fn=classify_installation, inputs=[uploaded_image], outputs=output_text)
-
-demo.launch()
+import os
+import torch
+import numpy as np
+import gradio as gr
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel
+
+# Load model parameters
+params = torch.load("clip_classification_params.pth", weights_only=False)
+
+correct_centroid = torch.tensor(params["correct_centroid"])
+incorrect_centroid = torch.tensor(params["incorrect_centroid"])
+threshold = params["threshold"]
+
+# Load CLIP model
+MODEL_NAME = "openai/clip-vit-base-patch32"
+clip_model = CLIPModel.from_pretrained(MODEL_NAME)
+clip_processor = CLIPProcessor.from_pretrained(MODEL_NAME)
+clip_model.eval()
+
+# Paths to example images
+correct_examples = [
+    "dataset/correct/correct (1).webp",
+    "dataset/correct/correct (2).webp",
+    "dataset/correct/correct (3).webp",
+    "dataset/correct/correct (4).webp"
+]
+
+incorrect_examples = [
+    "dataset/incorrect/incorrect (1).webp",
+    "dataset/incorrect/incorrect (2).webp",
+    "dataset/incorrect/incorrect (3).webp",
+     "dataset/incorrect/incorrect (4).webp"
+]
+
+# Function to classify an image
+def classify_installation(image):
+    """Classify if the bed installation is correct or incorrect."""
+    inputs = clip_processor(images=image, return_tensors="pt", padding=True)
+    with torch.no_grad():
+        embedding = clip_model.get_image_features(**inputs)
+        embedding = embedding / embedding.norm(p=2, dim=-1, keepdim=True)
+
+    # Compute similarity to correct centroid
+    similarity = torch.matmul(embedding, correct_centroid)
+    
+    # Compare similarity with threshold
+    if similarity.item() >= threshold:
+        return f"✅ Correct Installation (similarity = {similarity.item():.2f})"
+    else:
+        return f"❌ Incorrect Installation (similarity = {similarity.item():.2f})"
+
+# Function to load image from file path
+def load_image(image_path):
+    return Image.open(image_path).convert("RGB")
+
+# Function to process selected example image
+def process_example(image_path):
+    image = load_image(image_path)
+    return classify_installation(image), image
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# 🛏️ Bed Installation Classifier")
+    gr.Markdown("Upload an image or select one from the examples below to check if the bed installation is correct.")
+
+    with gr.Row():
+        uploaded_image = gr.Image(type="pil", label="Upload an image for testing")
+        output_text = gr.Textbox(label="Result")
+
+    gr.Markdown("### Check Installations (Click Button to Classify)")
+    with gr.Row():
+        correct_buttons = []
+        for i, img_path in enumerate(correct_examples):
+            with gr.Column():
+                gr.Image(value=load_image(img_path), interactive=False, width=150, height=150)
+                btn = gr.Button(value=f"Check Accuracy")
+                correct_buttons.append((btn, img_path))
+
+    with gr.Row():
+        incorrect_buttons = []
+        for i, img_path in enumerate(incorrect_examples):
+            with gr.Column():
+                gr.Image(value=load_image(img_path), interactive=False, width=150, height=150)
+                btn = gr.Button(value=f"Check Accuracy")
+                incorrect_buttons.append((btn, img_path))
+
+    # Connect buttons to classification function
+    for btn, img_path in correct_buttons + incorrect_buttons:
+        btn.click(fn=process_example, inputs=[gr.State(img_path)], outputs=[output_text, uploaded_image])
+
+    # Process uploaded image
+    uploaded_image.change(fn=classify_installation, inputs=[uploaded_image], outputs=output_text)
+
+demo.launch()