Spaces:
Sleeping
Sleeping
nazlicanto
commited on
Commit
•
6e3cb6a
1
Parent(s):
c2f84d5
Update app06.py
Browse files
app06.py
CHANGED
@@ -1,25 +1,17 @@
|
|
1 |
import streamlit as st
|
2 |
-
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
|
3 |
from PIL import Image
|
4 |
import numpy as np
|
5 |
import torch
|
6 |
import os
|
7 |
|
8 |
-
|
9 |
-
|
10 |
-
model = SegformerForSemanticSegmentation.from_pretrained(model_dir)
|
11 |
-
processor = SegformerImageProcessor.from_pretrained(model_dir)
|
12 |
-
model.eval()
|
13 |
|
|
|
|
|
14 |
|
15 |
-
|
16 |
-
model = SegformerForSemanticSegmentation.from_pretrained(model_dir, local_files_only=True)
|
17 |
-
preprocessor = SegformerImageProcessor.from_pretrained(model_dir, local_files_only=True)
|
18 |
-
except Exception as e:
|
19 |
-
existing_files = os.listdir(model_dir)
|
20 |
-
error_msg = f"Error: {e}\nFiles in directory: {existing_files}"
|
21 |
-
st.write(error_msg)
|
22 |
-
|
23 |
|
24 |
st.title("PCB Defect Detection")
|
25 |
|
@@ -29,14 +21,14 @@ uploaded_file = st.file_uploader("Upload a PCB image", type=["jpg", "png"])
|
|
29 |
if uploaded_file:
|
30 |
# Preprocess the image
|
31 |
test_image = Image.open(uploaded_file).convert("RGB")
|
32 |
-
inputs =
|
33 |
|
34 |
# Model inference
|
35 |
with torch.no_grad():
|
36 |
outputs = model(**inputs)
|
37 |
|
38 |
# Post-process
|
39 |
-
semantic_map =
|
40 |
semantic_map = np.uint8(semantic_map)
|
41 |
semantic_map[semantic_map==1] = 255
|
42 |
semantic_map[semantic_map==2] = 195
|
|
|
1 |
import streamlit as st
|
2 |
+
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor, SegformerConfig
|
3 |
from PIL import Image
|
4 |
import numpy as np
|
5 |
import torch
|
6 |
import os
|
7 |
|
8 |
+
model_path = "/home/user/app/defectdetection/model"
|
9 |
+
config = SegformerConfig.from_json_file(os.path.join(model_path, "config.json"))
|
|
|
|
|
|
|
10 |
|
11 |
+
model = SegformerForSemanticSegmentation(config=config)
|
12 |
+
model.load_state_dict(torch.load(os.path.join(model_path, "pytorch_model.bin")))
|
13 |
|
14 |
+
preprocessor = SegformerImageProcessor.from_pretrained(model_path, local_files_only=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
15 |
|
16 |
st.title("PCB Defect Detection")
|
17 |
|
|
|
21 |
if uploaded_file:
|
22 |
# Preprocess the image
|
23 |
test_image = Image.open(uploaded_file).convert("RGB")
|
24 |
+
inputs = preprocessor(images=test_image, return_tensors="pt")
|
25 |
|
26 |
# Model inference
|
27 |
with torch.no_grad():
|
28 |
outputs = model(**inputs)
|
29 |
|
30 |
# Post-process
|
31 |
+
semantic_map = preprocessor.post_process_semantic_segmentation(outputs, target_sizes=[test_image.size[::-1]])[0]
|
32 |
semantic_map = np.uint8(semantic_map)
|
33 |
semantic_map[semantic_map==1] = 255
|
34 |
semantic_map[semantic_map==2] = 195
|