deep-learning-analytics commited on
Commit
4626ab4
1 Parent(s): 432e4a1

added spinner for inference

Browse files
Files changed (1) hide show
  1. app.py +4 -2
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
- outputs = model(pixel_values=pixel_values)# logits are of shape (batch_size, num_labels, height/4, width/4)
 
 
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,