Update app06.py

app06.py

@@ -1,38 +1,38 @@
-
-import streamlit as st
-from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
-from PIL import Image
-import numpy as np
-import torch
-
-# Load the model and processor
-model_dir = "defectdetection/model/"
-model = SegformerForSemanticSegmentation.from_pretrained(model_dir)
-processor = SegformerImageProcessor.from_pretrained(model_dir)
-model.eval()
-
-st.title("PCB Defect Detection")
-
-# Upload image in Streamlit
-uploaded_file = st.file_uploader("Upload a PCB image", type=["jpg", "png"])
-
-if uploaded_file:
-    # Preprocess the image
-    test_image = Image.open(uploaded_file).convert("RGB")
-    inputs = processor(images=test_image, return_tensors="pt")
-
-    # Model inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    # Post-process
-    semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[test_image.size[::-1]])[0]
-    semantic_map = np.uint8(semantic_map)
-    semantic_map[semantic_map==1] = 255
-    semantic_map[semantic_map==2] = 195
-    semantic_map[semantic_map==3] = 135
-    semantic_map[semantic_map==4] = 75
-
-    # Display the results
-    st.image(test_image, caption="Uploaded Image", use_column_width=True)
-    st.image(semantic_map, caption="Predicted Defects", use_column_width=True, channels="GRAY")
+
+import streamlit as st
+from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
+from PIL import Image
+import numpy as np
+import torch
+
+# Load the model and processor
+model_dir = "defectdetection/model"
+model = SegformerForSemanticSegmentation.from_pretrained(model_dir)
+processor = SegformerImageProcessor.from_pretrained(model_dir)
+model.eval()
+
+st.title("PCB Defect Detection")
+
+# Upload image in Streamlit
+uploaded_file = st.file_uploader("Upload a PCB image", type=["jpg", "png"])
+
+if uploaded_file:
+    # Preprocess the image
+    test_image = Image.open(uploaded_file).convert("RGB")
+    inputs = processor(images=test_image, return_tensors="pt")
+
+    # Model inference
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    # Post-process
+    semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[test_image.size[::-1]])[0]
+    semantic_map = np.uint8(semantic_map)
+    semantic_map[semantic_map==1] = 255
+    semantic_map[semantic_map==2] = 195
+    semantic_map[semantic_map==3] = 135
+    semantic_map[semantic_map==4] = 75
+
+    # Display the results
+    st.image(test_image, caption="Uploaded Image", use_column_width=True)
+    st.image(semantic_map, caption="Predicted Defects", use_column_width=True, channels="GRAY")