Colab update - fix classify

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
+
+
+# This should be the same as the first line of Python code in this Colab notebook
+extractor = SegformerFeatureExtractor.from_pretrained("neerajprad/street-segmentation-demo")
+model = SegformerForSemanticSegmentation.from_pretrained("neerajprad/street-segmentation-demo")
+        
+def classify(im):
+  inputs = extractor(images=im, return_tensors="pt")
+  outputs = model(**inputs)
+  logits = outputs.logits
+  # Reshape output to be the same size as input
+  logits_tensor = nn.functional.interpolate(
+      logits,
+      size=im.shape[:2],
+      mode="bilinear",
+      align_corners=False,
+  )
+  classes = logits_tensor[0].detach().cpu().numpy().argmax(axis=0)
+  colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], 	[128,0,128], [0, 0, 0]])
+  return colors[classes]
+
+interface = gr.Interface(
+    title="Road segmentation",
+    description="Segment the scene into road/sidewalk, human, vehicles, other obj, nature/background",
+    examples=["ex1.jpg", "ex2.jpg", "ex3.jpg"],
+    fn=classify,
+    inputs=gr.Image(),
+    outputs=gr.Image(),
+)
+
+interface.launch(debug=True)

requirements.txt

@@ -0,0 +1,2 @@
+torch
+transformers