Create app.py

app.py

@@ -0,0 +1,36 @@
+import numpy as np
+import gradio as gr
+import torch
+from torch import nn
+from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor
+
+#extractor = AutoFeatureExtractor.from_pretrained("andresgtn/segformer-b0-finetuned-ade-64-64-finetuned-semantic-sidewalk")
+extractor = SegformerFeatureExtractor() 
+model = SegformerForSemanticSegmentation.from_pretrained("andresgtn/segformer-b0-finetuned-ade-64-64-finetuned-semantic-sidewalk")
+
+def rescale_output_image(logits, image):
+
+  upsampled_logits = nn.functional.interpolate(
+    logits,
+    size=image.shape[::-1][1:][::-1], # (height, width)
+    mode='bilinear',
+    align_corners=False
+    )
+  pred_seg = upsampled_logits.argmax(dim=1)[0]
+  return pred_seg
+  
+# classify function
+def classify(im):
+  inputs = extractor(images=im, return_tensors="pt")#.to("cuda")
+  outputs = model(**inputs)
+  logits = outputs.logits
+  #classes = logits[0].detach().cpu().numpy().argmax(axis=0)
+  #classes = rescale_output_image(logits, im).detach().cpu().numpy()
+  classes = rescale_output_image(logits, im).detach().numpy()
+  colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], 	[128,0,128], [0, 0, 0]])
+  return colors[classes]
+
+# define gradio interface 
+interface = gr.Interface(classify, gr.Image(), 'image')# FILL HERE
+
+interface.launch(debug=True)