Update app06.py

app06.py

@@ -1,25 +1,17 @@
 import streamlit as st
-from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
+from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor, SegformerConfig
 from PIL import Image
 import numpy as np
 import torch
 import os
 
-# Load the model and processor
-model_dir = "/home/user/app/defectdetection/model/"
-model = SegformerForSemanticSegmentation.from_pretrained(model_dir)
-processor = SegformerImageProcessor.from_pretrained(model_dir)
-model.eval()
+model_path = "/home/user/app/defectdetection/model"
+config = SegformerConfig.from_json_file(os.path.join(model_path, "config.json"))
 
+model = SegformerForSemanticSegmentation(config=config)
+model.load_state_dict(torch.load(os.path.join(model_path, "pytorch_model.bin")))
 
-try:
-    model = SegformerForSemanticSegmentation.from_pretrained(model_dir, local_files_only=True)
-    preprocessor = SegformerImageProcessor.from_pretrained(model_dir, local_files_only=True)
-except Exception as e:
-    existing_files = os.listdir(model_dir)
-    error_msg = f"Error: {e}\nFiles in directory: {existing_files}"
-    st.write(error_msg)
-
+preprocessor = SegformerImageProcessor.from_pretrained(model_path, local_files_only=True)
 
 st.title("PCB Defect Detection")
 
@@ -29,14 +21,14 @@ 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")
+    inputs = preprocessor(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 = preprocessor.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