hshetty commited on
Commit
c1d11dd
1 Parent(s): a8d2fde

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -6,12 +6,15 @@ import numpy as np
6
 
7
  extractor = AutoFeatureExtractor.from_pretrained("hshetty/segmentation-model-finetuned-on-semantic-sidewalk-3e-4-e5")
8
  model = SegformerForSemanticSegmentation.from_pretrained("hshetty/segmentation-model-finetuned-on-semantic-sidewalk-3e-4-e5")
 
9
  def classify(im):
10
  inputs = extractor(images=im, return_tensors="pt")
11
  outputs = model(**inputs)
12
  logits = outputs.logits
13
  classes = logits[0].detach().cpu().numpy().argmax(axis=0)
14
  colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], [128,0,128], [0, 0, 0]])
 
 
15
  return colors[classes]
16
 
17
 
 
6
 
7
  extractor = AutoFeatureExtractor.from_pretrained("hshetty/segmentation-model-finetuned-on-semantic-sidewalk-3e-4-e5")
8
  model = SegformerForSemanticSegmentation.from_pretrained("hshetty/segmentation-model-finetuned-on-semantic-sidewalk-3e-4-e5")
9
+ counter = 0
10
  def classify(im):
11
  inputs = extractor(images=im, return_tensors="pt")
12
  outputs = model(**inputs)
13
  logits = outputs.logits
14
  classes = logits[0].detach().cpu().numpy().argmax(axis=0)
15
  colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], [128,0,128], [0, 0, 0]])
16
+ counter+=1
17
+ print("hshetty:{}".format(counter))
18
  return colors[classes]
19
 
20