Updated from colab

app.py

@@ -7,8 +7,6 @@ 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)