Commit
·
4626ab4
1
Parent(s):
432e4a1
added spinner for inference
Browse files
app.py
CHANGED
@@ -2,12 +2,12 @@ import os
|
|
2 |
import pandas as pd
|
3 |
import numpy as np
|
4 |
import torch
|
5 |
-
from PIL import Image
|
6 |
from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor
|
7 |
from torch import nn
|
8 |
import streamlit as st
|
9 |
|
10 |
|
|
|
11 |
raw_image = st.file_uploader('Raw Input Image')
|
12 |
if raw_image is not None:
|
13 |
df = pd.read_csv('class_dict_seg.csv')
|
@@ -30,7 +30,9 @@ if raw_image is not None:
|
|
30 |
feature_extractor_inference = SegformerFeatureExtractor(do_random_crop=False, do_pad=False)
|
31 |
pixel_values = feature_extractor_inference(image, return_tensors="pt").pixel_values.to(device)
|
32 |
model.eval()
|
33 |
-
|
|
|
|
|
34 |
logits = outputs.logits.cpu()
|
35 |
# First, rescale logits to original image size
|
36 |
upsampled_logits = nn.functional.interpolate(logits,
|
|
|
2 |
import pandas as pd
|
3 |
import numpy as np
|
4 |
import torch
|
|
|
5 |
from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor
|
6 |
from torch import nn
|
7 |
import streamlit as st
|
8 |
|
9 |
|
10 |
+
st.title('Semantic Segmentation using SegFormer')
|
11 |
raw_image = st.file_uploader('Raw Input Image')
|
12 |
if raw_image is not None:
|
13 |
df = pd.read_csv('class_dict_seg.csv')
|
|
|
30 |
feature_extractor_inference = SegformerFeatureExtractor(do_random_crop=False, do_pad=False)
|
31 |
pixel_values = feature_extractor_inference(image, return_tensors="pt").pixel_values.to(device)
|
32 |
model.eval()
|
33 |
+
|
34 |
+
with st.spinner('Running inference...'):
|
35 |
+
outputs = model(pixel_values=pixel_values)# logits are of shape (batch_size, num_labels, height/4, width/4)
|
36 |
logits = outputs.logits.cpu()
|
37 |
# First, rescale logits to original image size
|
38 |
upsampled_logits = nn.functional.interpolate(logits,
|