nazlicanto commited on
Commit
355d137
2 Parent(s): 046c340 36ac725

model_dir update w/ absolute path

Browse files
Files changed (1) hide show
  1. app06.py +82 -41
app06.py CHANGED
@@ -1,41 +1,82 @@
1
-
2
- import streamlit as st
3
- from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
4
- from PIL import Image
5
- import numpy as np
6
- import torch
7
-
8
- # Load the model and processor
9
- model_dir = "/home/user/app/defectdetection/model"
10
- model = SegformerForSemanticSegmentation.from_pretrained(model_path)
11
- preprocessor = SegformerImageProcessor.from_pretrained(model_path)
12
-
13
- model = SegformerForSemanticSegmentation.from_pretrained(model_dir)
14
- processor = SegformerImageProcessor.from_pretrained(model_dir)
15
- model.eval()
16
-
17
- st.title("PCB Defect Detection")
18
-
19
- # Upload image in Streamlit
20
- uploaded_file = st.file_uploader("Upload a PCB image", type=["jpg", "png"])
21
-
22
- if uploaded_file:
23
- # Preprocess the image
24
- test_image = Image.open(uploaded_file).convert("RGB")
25
- inputs = processor(images=test_image, return_tensors="pt")
26
-
27
- # Model inference
28
- with torch.no_grad():
29
- outputs = model(**inputs)
30
-
31
- # Post-process
32
- semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[test_image.size[::-1]])[0]
33
- semantic_map = np.uint8(semantic_map)
34
- semantic_map[semantic_map==1] = 255
35
- semantic_map[semantic_map==2] = 195
36
- semantic_map[semantic_map==3] = 135
37
- semantic_map[semantic_map==4] = 75
38
-
39
- # Display the results
40
- st.image(test_image, caption="Uploaded Image", use_column_width=True)
41
- st.image(semantic_map, caption="Predicted Defects", use_column_width=True, channels="GRAY")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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