Change

app.py

@@ -32,17 +32,8 @@ def preprocess(image):
 def predict(img):
     img = preprocess(img)
     preds = model.predict(img)[0]
-    cn = ""
-    if class_names[0] in ["apples", "bananas", "eggshells", "generalcompost", "peels"]:
-        cn = "Food Waste"
-    elif class_names[0] in ["bottles", "cans", "platicbags", "plastics"]:
-        cn = "Plastic Waste"
-    elif class_names[0] in ["tissue papers", "cardboard"]:
-        cn = "Paper waste"
-    else:
-        cn = "Others"
-        
-    return {cn: float(preds[0])}
+    
+    return {class_names[i]: float(preds[i]) for i in range(len(preds))}
 
 demo = gr.Interface(
     fn=predict,