Update app.py

app.py

@@ -8,11 +8,17 @@ import matplotlib.pyplot as plt
 extractor = AutoFeatureExtractor.from_pretrained("brendenc/my-segmentation-model")
 model = AutoModelForImageClassification.from_pretrained("brendenc/my-segmentation-model")
 
+collapse_categories = {**{i: 0 for i in range(1, 8)}, 
+                            **{i: 1 for i in range(8, 10)}, 
+                            **{i: 2 for i in range(10, 18)}, 
+                            **{i: 3 for i in range(18, 28)}}
+                            
 def classify(im):
   inputs = extractor(images=im, return_tensors="pt")
   outputs = model(**inputs)
   logits = outputs.logits
   classes = logits[0].detach().numpy().argmax(axis=0)
+  classes = np.vectorize(lambda x: collapse_categories.get(x, 4))(classes)
   colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], 	[128,0,128], [0, 0, 0]])
   return colors[classes]