Spaces:
Sleeping
Sleeping
File size: 1,477 Bytes
b47f30d 6e3cb6a b47f30d cfd86c0 b47f30d 6e3cb6a 355d137 6e3cb6a cfd86c0 6e3cb6a cfd86c0 355d137 6e3cb6a 355d137 6e3cb6a 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 |
import streamlit as st
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor, SegformerConfig
from PIL import Image
import numpy as np
import torch
import os
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")))
preprocessor = SegformerImageProcessor.from_pretrained(model_path, local_files_only=True)
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 = preprocessor(images=test_image, return_tensors="pt")
# Model inference
with torch.no_grad():
outputs = model(**inputs)
# Post-process
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
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")
|