Update app.py

app.py

@@ -10,8 +10,7 @@ import colorsys
 learn = load_learner('fashion_mnist_resnet18.pkl')
 
 # Define class names
-class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
-               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
+class_names = ['jeans', 'top']
 
 # Helper function to get the name of the closest color
 def closest_color(requested_color):
@@ -41,18 +40,12 @@ def get_complementary_color(rgb_color):
     return closest_color((r, g, b))
 
 def get_outfit_recommendation(pred_class):
-    if pred_class in ['T-shirt/top', 'Pullover', 'Shirt']:
-        return 'Trousers'
-    elif pred_class == 'Trouser':
-        return 'Top'
-    elif pred_class in ['Dress', 'Coat']:
-        return 'Shoes'
-    elif pred_class in ['Sandal', 'Sneaker', 'Ankle boot']:
-        return 'Trousers'
-    elif pred_class == 'Bag':
-        return 'Top'
+    if pred_class in ['jeans']:
+        return 'top'
+    elif pred_class in ['top']:
+        return 'top'
     else:
-        return 'Item'
+        return 'item'
 
 def predict(image):
     pred_class, pred_idx, outputs = learn.predict(image)