Update app.py

app.py

@@ -22,14 +22,15 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 model.eval()
 
-# Define the image transformations (adjust as needed for your model)
+# Define the image transformations
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
 ])
 
-# Define class labels
+# Define class labels (German)
 class_labels = [
     "Spraydosen", "Lebensmitteldosen_aus_Aluminium", "Getränkedosen_aus_Aluminium",
     "Pappkartons", "Verpackungen_aus_Pappe", "Kleidung", "Kaffeekränzchen",
@@ -43,9 +44,10 @@ class_labels = [
     "Teebeutel"
 ]
 
+# Mapping classes to correct trash bin
 class_to_tonne = {
     "Spraydosen": "Gelbe Tonne",
-    "Lebensmitteldosen_aus_Aluminiumm": "Gelbe Tonne",
+    "Lebensmitteldosen_aus_Aluminium": "Gelbe Tonne",
     "Getränkedosen_aus_Aluminium": "Gelbe Tonne",
     "Plastiktüten": "Gelbe Tonne",
     "Wasserflaschen_aus_Kunststoff": "Gelbe Tonne",
@@ -67,7 +69,7 @@ class_to_tonne = {
     "Teebeutel": "Biomüll",
     "Getränkeflaschen_aus_Glas": "Altglas",
     "Lebensmittelgläser_aus_Glas": "Altglas",
-    "Losmetikbehälter_aus_Glas": "Altglas",
+    "Kosmetikbehälter_aus_Glas": "Altglas",
     "Kleidung": "Restmüll",
     "Schuhe": "Restmüll",
     "Styroporbecher": "Restmüll",
@@ -75,22 +77,10 @@ class_to_tonne = {
     "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(model.device)
-    with torch.no_grad():
-        outputs = model(input_tensor)
-        _, predicted = torch.max(outputs, 1)
-        label = class_labels[predicted.item()]
-        tonne = class_to_tonne.get(label, "Unbekannt")
-    return label.replace("_", " ").capitalize(), tonne
-
-# Gradio interface setup
+# Gradio Blocks UI
 with gr.Blocks(title="Abfallerkennung mit KI 🗑️") as demo:
     gr.Markdown("### Lade ein Bild hoch und erfahre, in welche Tonne der Abfall gehört.")
+    
     image_input = gr.Image(type="pil", label="Bild hochladen")
     output_label = gr.Textbox(label="Vorhersage")
     output_tonne = gr.Textbox(label="Richtiger Abfallbehälter")
@@ -99,7 +89,7 @@ with gr.Blocks(title="Abfallerkennung mit KI 🗑️") as demo:
     def analyze(image):
         if image.mode != "RGB":
             image = image.convert("RGB")
-        input_tensor = transform(image).unsqueeze(0).to(model.device)
+        input_tensor = transform(image).unsqueeze(0).to(device)
         with torch.no_grad():
             outputs = model(input_tensor)
             _, predicted = torch.max(outputs, 1)