Spaces:
Sleeping
Sleeping
nazlicanto
commited on
model_dir update w/ absolute path
Browse files
app06.py
CHANGED
@@ -1,41 +1,82 @@
|
|
1 |
-
|
2 |
-
|
3 |
-
|
4 |
-
from
|
5 |
-
|
6 |
-
import
|
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 |
-
semantic_map =
|
34 |
-
semantic_map
|
35 |
-
semantic_map[semantic_map==
|
36 |
-
semantic_map[semantic_map==
|
37 |
-
semantic_map[semantic_map==
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
st.image(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
<<<<<<< HEAD
|
2 |
+
|
3 |
+
import streamlit as st
|
4 |
+
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
|
5 |
+
from PIL import Image
|
6 |
+
import numpy as np
|
7 |
+
import torch
|
8 |
+
|
9 |
+
# Load the model and processor
|
10 |
+
model_dir = "/home/user/app/defectdetection/model"
|
11 |
+
model = SegformerForSemanticSegmentation.from_pretrained(model_path)
|
12 |
+
preprocessor = SegformerImageProcessor.from_pretrained(model_path)
|
13 |
+
|
14 |
+
model = SegformerForSemanticSegmentation.from_pretrained(model_dir)
|
15 |
+
processor = SegformerImageProcessor.from_pretrained(model_dir)
|
16 |
+
model.eval()
|
17 |
+
|
18 |
+
st.title("PCB Defect Detection")
|
19 |
+
|
20 |
+
# Upload image in Streamlit
|
21 |
+
uploaded_file = st.file_uploader("Upload a PCB image", type=["jpg", "png"])
|
22 |
+
|
23 |
+
if uploaded_file:
|
24 |
+
# Preprocess the image
|
25 |
+
test_image = Image.open(uploaded_file).convert("RGB")
|
26 |
+
inputs = processor(images=test_image, return_tensors="pt")
|
27 |
+
|
28 |
+
# Model inference
|
29 |
+
with torch.no_grad():
|
30 |
+
outputs = model(**inputs)
|
31 |
+
|
32 |
+
# Post-process
|
33 |
+
semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[test_image.size[::-1]])[0]
|
34 |
+
semantic_map = np.uint8(semantic_map)
|
35 |
+
semantic_map[semantic_map==1] = 255
|
36 |
+
semantic_map[semantic_map==2] = 195
|
37 |
+
semantic_map[semantic_map==3] = 135
|
38 |
+
semantic_map[semantic_map==4] = 75
|
39 |
+
|
40 |
+
# Display the results
|
41 |
+
st.image(test_image, caption="Uploaded Image", use_column_width=True)
|
42 |
+
st.image(semantic_map, caption="Predicted Defects", use_column_width=True, channels="GRAY")
|
43 |
+
=======
|
44 |
+
|
45 |
+
import streamlit as st
|
46 |
+
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
|
47 |
+
from PIL import Image
|
48 |
+
import numpy as np
|
49 |
+
import torch
|
50 |
+
|
51 |
+
# Load the model and processor
|
52 |
+
model_dir = "/home/user/app/defectdetection/model"
|
53 |
+
model = SegformerForSemanticSegmentation.from_pretrained(model_dir)
|
54 |
+
processor = SegformerImageProcessor.from_pretrained(model_dir)
|
55 |
+
model.eval()
|
56 |
+
|
57 |
+
st.title("PCB Defect Detection")
|
58 |
+
|
59 |
+
# Upload image in Streamlit
|
60 |
+
uploaded_file = st.file_uploader("Upload a PCB image", type=["jpg", "png"])
|
61 |
+
|
62 |
+
if uploaded_file:
|
63 |
+
# Preprocess the image
|
64 |
+
test_image = Image.open(uploaded_file).convert("RGB")
|
65 |
+
inputs = processor(images=test_image, return_tensors="pt")
|
66 |
+
|
67 |
+
# Model inference
|
68 |
+
with torch.no_grad():
|
69 |
+
outputs = model(**inputs)
|
70 |
+
|
71 |
+
# Post-process
|
72 |
+
semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[test_image.size[::-1]])[0]
|
73 |
+
semantic_map = np.uint8(semantic_map)
|
74 |
+
semantic_map[semantic_map==1] = 255
|
75 |
+
semantic_map[semantic_map==2] = 195
|
76 |
+
semantic_map[semantic_map==3] = 135
|
77 |
+
semantic_map[semantic_map==4] = 75
|
78 |
+
|
79 |
+
# Display the results
|
80 |
+
st.image(test_image, caption="Uploaded Image", use_column_width=True)
|
81 |
+
st.image(semantic_map, caption="Predicted Defects", use_column_width=True, channels="GRAY")
|
82 |
+
>>>>>>> 36ac725dd03eaeedd3c4601d12a8b80b846b7647
|