Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import ViTForImageClassification, ViTFeatureExtractor
 from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 from PIL import Image
 
-# Load crop model (ViT)
+# Load crop disease model (ViT)
 model = ViTForImageClassification.from_pretrained("iamomtiwari/VITPEST")
 feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")
 
@@ -12,6 +12,10 @@ feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")
 fallback_model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
 fallback_feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
 
+# Load additional ViT model (221k model) for a different classification if the user feedback is "no"
+vit_221k_model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224-in21k")
+vit_221k_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
+
 # Define class labels with treatment advice
 class_labels = {
     1: {"label": "Stage Corn___Common_Rust", "treatment": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants."},
@@ -34,7 +38,7 @@ class_labels = {
 labels_list = [class_labels[i]["label"] for i in range(1, 15)]
 
 # Inference function
-def predict(image):
+def predict(image, feedback):
     # First, use the crop disease model (ViT)
     inputs = feature_extractor(images=image, return_tensors="pt")
     with torch.no_grad():
@@ -44,10 +48,22 @@ def predict(image):
     # Check if the predicted label corresponds to a crop disease
     predicted_label = labels_list[predicted_class_idx]
 
-    # Ask user if the prediction is correct
+    # If the feedback is "no", switch to ViT 221k model for a different class prediction
+    if feedback == "no":
+        # Use ViT 221k model
+        inputs_vit_221k = vit_221k_feature_extractor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            outputs_vit_221k = vit_221k_model(**inputs_vit_221k)
+            predicted_class_idx_vit_221k = outputs_vit_221k.logits.argmax(-1).item()
+        
+        # Get the ViT 221k prediction label
+        vit_221k_label = vit_221k_model.config.id2label[predicted_class_idx_vit_221k]
+        return f"Fallback ViT 221k Prediction: {vit_221k_label}"
+
+    # If feedback is "yes", return the initial disease prediction and treatment advice
     if predicted_class_idx < len(class_labels):  # It's a crop disease
         treatment_advice = class_labels[predicted_class_idx + 1]["treatment"]
-        return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}\n\nIs this prediction correct? (yes/no)"
+        return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}"
     else:
         # If not a crop disease, use the fallback model (ResNet50) for general object detection
         inputs_fallback = fallback_feature_extractor(images=image, return_tensors="pt")
@@ -57,8 +73,12 @@ def predict(image):
 
         # Get the fallback prediction label
         fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]
-        return f"Fallback Prediction (Not a Crop): {fallback_label}\n\nIs this prediction correct? (yes/no)"
+        return f"Fallback Prediction (Not a Crop): {fallback_label}"
 
 # Create Gradio Interface
-interface = gr.Interface(fn=predict, inputs="image", outputs="text")
+interface = gr.Interface(
+    fn=predict, 
+    inputs=["image", gr.Radio(["yes", "no"], label="Is the prediction correct?")],
+    outputs="text"
+)
 interface.launch()