Add chat model

app.py

@@ -5,31 +5,29 @@ import torch
 from transformers import T5Tokenizer, T5ForConditionalGeneration  #for chat model
 
 # ------------------ Load classifier model + extractor ------------------
-model_name = "Aalaa/Fine_tuned_Vit_trash_classification"
+cls_model_name = "Aalaa/Fine_tuned_Vit_trash_classification"
+feature_extractor = AutoFeatureExtractor.from_pretrained(cls_model_name)
+cls_model = AutoModelForImageClassification.from_pretrained(cls_model_name)
 
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
-# Label mapping
-id2label = model.config.id2label
+id2label = cls_model.config.id2label
 
-def predict(image):
-    # Convert Gradio PIL to model input
+
+def classify_image(image):
+    """Run ViT classifier and return top 3 predictions."""
     inputs = feature_extractor(images=image, return_tensors="pt")
 
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = cls_model(**inputs)
 
     logits = outputs.logits
     probs = torch.nn.functional.softmax(logits, dim=-1)[0]
 
-    # Return top 3
-    result = {
-        id2label[i]: float(probs[i])
-        for i in probs.topk(3).indices.tolist()
-    }
+    top3_indices = probs.topk(3).indices.tolist()
 
+    result = {id2label[i]: float(probs[i]) for i in top3_indices}
     return result
 
+
 #------------------ Load flan-t5 chat model------------------
 tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
 chat_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
@@ -60,11 +58,33 @@ def explain_recycling(class_label):
 
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=3),
-    title="AI Waste Classifier"
-)
+# ------------------ GRADIO APP ------------------
+
+def full_pipeline(image):
+    """Classifier → get top label → Flan-T5 explanation."""
+    predictions = classify_image(image)
+
+    # Get the top class name
+    top_label = max(predictions, key=predictions.get)
+
+    explanation = explain_recycling(top_label)
+
+    return predictions, explanation
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# ♻️ AI Waste Classifier + Disposal Advisor")
+
+    with gr.Row():
+        img_input = gr.Image(type="pil", label="Upload waste image")
+
+    with gr.Column():
+        cls_output = gr.Label(num_top_classes=3, label="Classifier Prediction")
+        explain_output = gr.Textbox(label="Recycling Advice", lines=6)
+
+    run_btn = gr.Button("Analyze")
+
+    run_btn.click(full_pipeline, inputs=img_input, outputs=[cls_output, explain_output])
+
 
 demo.launch()