Updated from colab

app.py

@@ -0,0 +1,28 @@
+import gradio as gr
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+import torch
+import matplotlib.pyplot as plt
+
+
+feature_extractor = AutoFeatureExtractor.from_pretrained("brendenc/my-segmentation-model")
+model = AutoModelForImageClassification.from_pretrained("brendenc/my-segmentation-model")
+
+labels = dataset['train'].features['labels'].names
+
+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)
+  colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], 	[128,0,128], [0, 0, 0]])
+  return colors[classes]
+
+example_imgs = [f"example_{i}.jpg" for i in range(3)]
+interface = gr.Interface(classify,
+                         inputs="image",
+                         outputs="image",
+                         examples = example_imgs,
+                         title = "Street Image Segmentation",
+                         description = """Below is a simple app for image segmentation. This model was trained using""")
+
+interface.launch(debug=True)

requirements.txt

@@ -0,0 +1,3 @@
+torch
+transformers
+matplotlib