Update app.py

app.py

@@ -17,19 +17,16 @@ logger = logging.getLogger(__name__)
 # Define the number of classes
 num_classes = 3
 
-# Confidence threshold for reliable predictions
+# Confidence threshold for main model predictions
 CONFIDENCE_THRESHOLD = 0.8  # 80%
 
-# Entropy threshold for flat probability distribution (to detect non-maize/rice images)
-ENTROPY_THRESHOLD = 0.9  # Lower entropy means a more peaked distribution (more confident)
-
 # Download model from Hugging Face
 def download_model():
     model_path = hf_hub_download(repo_id="jays009/Resnet3", filename="model.pth")
     return model_path
 
-# Load the model from Hugging Face
-def load_model(model_path):
+# Load the main model from Hugging Face
+def load_main_model(model_path):
     model = models.resnet50(pretrained=False)
     num_features = model.fc.in_features
     model.fc = nn.Sequential(
@@ -55,7 +52,7 @@ def load_model(model_path):
 
 # Path to your model
 model_path = download_model()
-model = load_model(model_path)
+main_model = load_main_model(model_path)
 
 # Define the transformation for the input image
 transform = transforms.Compose([
@@ -65,16 +62,6 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
-# Function to compute entropy of the probability distribution
-def compute_entropy(probabilities):
-    probs = probabilities.numpy()  # Convert to numpy array
-    # Avoid log(0) by adding a small epsilon
-    probs = np.clip(probs, 1e-10, 1.0)
-    entropy = -np.sum(probs * np.log(probs))
-    # Normalize entropy by the maximum possible entropy (log(num_classes))
-    max_entropy = np.log(num_classes)
-    return entropy / max_entropy
-
 # Prediction function for an uploaded image
 def predict_from_image_url(image_url):
     try:
@@ -89,9 +76,13 @@ def predict_from_image_url(image_url):
 
         # Perform prediction
         with torch.no_grad():
-            outputs = model(image_tensor)  # Shape: [1, 3]
+            outputs = main_model(image_tensor)  # Shape: [1, 3]
             logger.info(f"Model output shape: {outputs.shape}")
+            # Log raw logits
+            logger.info(f"Raw logits: {outputs[0].numpy()}")
             probabilities = torch.softmax(outputs, dim=1)[0]  # Convert to probabilities
+            # Log softmax probabilities
+            logger.info(f"Softmax probabilities: {probabilities.numpy()}")
             predicted_class = torch.argmax(outputs, dim=1).item()
 
         # Define class information
@@ -115,25 +106,7 @@ def predict_from_image_url(image_url):
         predicted_info = class_info[predicted_class]
         predicted_name = predicted_info["name"]
         problem_id = predicted_info["problem_id"]
-        confidence = probabilities[predicted_class].item()  # Confidence score for the predicted class
-
-        # Compute entropy of the probability distribution
-        entropy = compute_entropy(probabilities)
-        logger.info(f"Prediction entropy: {entropy:.4f}, confidence: {confidence:.4f} for image URL: {image_url}")
-
-        # Check if the image is likely maize or rice based on entropy and confidence
-        # High entropy (flat distribution) suggests the image may not be maize or rice
-        if entropy > ENTROPY_THRESHOLD:
-            logger.warning(
-                f"High entropy ({entropy:.4f} > {ENTROPY_THRESHOLD}) for image URL: {image_url}. "
-                "Image may not be of maize or rice."
-            )
-            return {
-                "status": "invalid",
-                "predicted_class": predicted_name,
-                "problem_id": problem_id,
-                "confidence": f"{confidence*100:.2f}%"
-            }
+        confidence = probabilities[predicted_class].item()
 
         # Check confidence threshold
         if confidence < CONFIDENCE_THRESHOLD: