Spaces:
Sleeping
Sleeping
File size: 1,673 Bytes
b47f30d cfd86c0 b47f30d c2f84d5 cfd86c0 355d137 cfd86c0 355d137 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import streamlit as st
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
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()
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)
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")
|