Update app.py

app.py

@@ -31,35 +31,83 @@ transform = transforms.Compose([
 
 # Define class labels
 class_labels = [
-    "aerosol_cans", "aluminum_food_cans", "aluminum_soda_cans", "cardboard_boxes", 
-    "cardboard_packaging", "clothing", "coffee_grounds", "disposable_plastic_cutlery", 
-    "eggshells", "food_waste", "glass_beverage_bottles", "glass_cosmetic_containers", 
-    "glass_food_jars", "magazines", "newspaper", "office_paper", "paper_cups", 
-    "plastic_cup_lids", "plastic_detergent_bottles", "plastic_food_containers", 
-    "plastic_shopping_bags", "plastic_soda_bottles", "plastic_straws", "plastic_trash_bags", 
-    "plastic_water_bottles", "shoes", "steel_food_cans", "styrofoam_cups", 
-    "styrofoam_food_containers", "tea_bags"
+    "Spraydosen", "Lebensmitteldosen_aus_Aluminium", "Getränkedosen_aus_Aluminium",
+    "Pappkartons", "Verpackungen_aus_Pappe", "Kleidung", "Kaffeekränzchen",
+    "Einweg_Besteck_aus_Kunststoff", "Eierschalen", "Lebensmittelabfälle",
+    "Getränkeflaschen_aus_Glas", "Kosmetikbehälter_aus_Glas", "Lebensmittelgläser_aus_Glas",
+    "Zeitschriften", "Zeitungen", "Buropapier", "Pappbecher", "Deckel_aus_Kunststoff",
+    "Waschmittelbehälter_aus_Kunststoff", "Lebensmittelbehälter_aus_Kunststoff",
+    "Plastiktüten", "Sodaflaschen_aus_Kunststoff", "Strohhalme_aus_Kunststoff",
+    "Mülltüten_aus_Kunststoff", "Wasserflaschen_aus_Kunststoff", "Schuhe",
+    "Lebensmitteldosen_aus_Stahl", "Styroporbecher", "Lebensmittelbehälter_aus_Styropor",
+    "Teebeutel"
 ]
 
+class_to_tonne = {
+    "Spraydosen": "Gelbe Tonne",
+    "Lebensmitteldosen_aus_Aluminiumm": "Gelbe Tonne",
+    "Getränkedosen_aus_Aluminium": "Gelbe Tonne",
+    "Plastiktüten": "Gelbe Tonne",
+    "Wasserflaschen_aus_Kunststoff": "Gelbe Tonne",
+    "Sodaflaschen_aus_Kunststoff": "Gelbe Tonne",
+    "Lebensmittelbehälter_aus_Kunststoff": "Gelbe Tonne",
+    "Waschmittelbehälter_aus_Kunststoff": "Gelbe Tonne",
+    "Deckel_aus_Kunststoff": "Gelbe Tonne",
+    "Strohhalme_aus_Kunststoff": "Gelbe Tonne",
+    "Einweg_Besteck_aus_Kunststoff": "Gelbe Tonne",
+
+    "Pappkartons": "Papiertonne",
+    "Verpackungen_aus_Pappe": "Papiertonne",
+    "Zeitschriften": "Papiertonne",
+    "Zeitungen": "Papiertonne",
+    "Buropapier": "Papiertonne",
+    "Pappbecher": "Papiertonne",
+
+    "Kaffeekränzchen": "Biomüll",
+    "Lebensmittelabfälle": "Biomüll",
+    "Eierschalen": "Biomüll",
+    "Teebeutel": "Biomüll",
+
+    "Getränkeflaschen_aus_Glas": "Altglas",
+    "Lebensmittelgläser_aus_Glas": "Altglas",
+    "Losmetikbehälter_aus_Glas": "Altglas",
+
+    "Kleidung": "Restmüll",
+    "Schuhe": "Restmüll",
+    "Styroporbecher": "Restmüll",
+    "Lebensmittelbehälter_aus_Styropor": "Restmüll",
+    "Mülltüten_aus_Kunststoff": "Restmüll"
+}
+
+
 # Prediction function
 def predict_image(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
-    input_tensor = transform(image).unsqueeze(0).to(device)
+    input_tensor = transform(image).unsqueeze(0).to(model.device)
     with torch.no_grad():
         outputs = model(input_tensor)
         _, predicted = torch.max(outputs, 1)
         label = class_labels[predicted.item()]
-    return label
+        tonne = class_to_tonne.get(label, "Unbekannt")
+    
+    return {
+        "Klasse (auf Deutsch)": label.replace("_", " ").capitalize(),
+        "Gehört in die": tonne
+    }
 
 # Gradio interface setup
 interface = gr.Interface(
     fn=predict_image,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs="text",
+    inputs=gr.Image(type="pil", label="Bild hochladen"),
+    outputs=[
+        gr.Textbox(label="Vorhersage (Klasse auf Deutsch)"),
+        gr.Textbox(label="Richtiger Abfallbehälter")
+    ],
+    title="Abfallerkennung mit KI 🗑️",
+    description="Lade ein Bild hoch und erfahre, in welche Tonne der Abfall gehört.",
     live=True
 )
 
-# Launch Gradio app
 interface.launch()