Create app.py

app.py

@@ -0,0 +1,157 @@
+import gradio as gr
+import torch
+import torchvision.models as models
+from torchvision.models import EfficientNet_B0_Weights # Or the specific version used
+from PIL import Image
+from torchvision import transforms
+import json
+from huggingface_hub import hf_hub_download
+import os
+
+# --- Configuration ---
+# This should be the ID of the repository where your MODEL is stored
+MODEL_REPO_ID = "bhumong/fruit-classifier-efficientnet-b0" # <-- REPLACE if different
+MODEL_FILENAME = "pytorch_model.bin"
+CONFIG_FILENAME = "config.json"
+
+# --- 1. Load Model and Config ---
+# (Using the function defined previously to load from Hub)
+def load_model_from_hf(repo_id, model_filename, config_filename):
+    """Loads model state_dict and config from Hugging Face Hub."""
+    try:
+        config_path = hf_hub_download(repo_id=repo_id, filename=config_filename)
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+        print("Config loaded:", config) # Debug print
+    except Exception as e:
+        print(f"Error loading config from {repo_id}/{config_filename}: {e}")
+        raise # Re-raise error if config fails
+
+    num_labels = config.get('num_labels')
+    id2label = config.get('id2label')
+
+    if num_labels is None or id2label is None:
+        raise ValueError("Config file must contain 'num_labels' and 'id2label'")
+
+    # Instantiate the correct architecture (EfficientNet-B0)
+    model = models.efficientnet_b0(weights=None) # Load architecture only
+
+    # Modify the classifier head
+    try:
+        num_ftrs = model.classifier[1].in_features
+        model.classifier[1] = torch.nn.Linear(num_ftrs, num_labels)
+    except Exception as e:
+        print(f"Error modifying model classifier: {e}")
+        raise
+
+    # Download and load model weights
+    try:
+        model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        state_dict = torch.load(model_path, map_location=device)
+        model.load_state_dict(state_dict)
+        model.to(device) # Move model to device
+        model.eval() # Set to evaluation mode
+        print(f"Model loaded successfully from {repo_id} to device {device}.")
+        return model, config, id2label, device
+    except Exception as e:
+        print(f"Error loading model weights from {repo_id}/{model_filename}: {e}")
+        raise
+
+# Load the model globally when the script starts
+try:
+    model, config, id2label, device = load_model_from_hf(MODEL_REPO_ID, MODEL_FILENAME, CONFIG_FILENAME)
+except Exception as e:
+    print(f"FATAL: Could not load model or config. Gradio app cannot start. Error: {e}")
+    # Optionally, exit or raise a specific error for Gradio to catch if possible
+    model, config, id2label, device = None, None, None, None # Prevent further errors
+
+
+# --- 2. Define Preprocessing ---
+IMG_SIZE = (224, 224)
+mean=[0.485, 0.456, 0.406]
+std=[0.229, 0.224, 0.225]
+
+preprocess = transforms.Compose([
+    transforms.Resize(IMG_SIZE),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=mean, std=std),
+])
+
+# --- 3. Define Prediction Function ---
+def predict(inp_image):
+    """Takes a PIL image, preprocesses, predicts, and returns label confidences."""
+    if model is None or id2label is None:
+         return {"Error": 1.0, "Message": "Model not loaded"} # Handle model load failure
+
+    if inp_image is None:
+        return {"Error": 1.0, "Message": "No image provided"}
+
+    try:
+        # Ensure image is RGB
+        img = inp_image.convert("RGB")
+        input_tensor = preprocess(img)
+        input_batch = input_tensor.unsqueeze(0) # Add batch dimension
+        input_batch = input_batch.to(device) # Move tensor to the correct device
+
+        with torch.no_grad():
+            output = model(input_batch)
+            probabilities = torch.nn.functional.softmax(output[0], dim=0)
+
+        # Prepare output for Gradio Label component (dictionary {label: probability})
+        confidences = {id2label[str(i)]: float(probabilities[i]) for i in range(len(id2label))}
+        return confidences
+
+    except Exception as e:
+        print(f"Error during prediction: {e}")
+        return {"Error": 1.0, "Message": f"Prediction failed: {e}"}
+
+
+# --- 4. Create Gradio Interface ---
+
+# Add example images (Make sure these paths exist within your Space repo!)
+# Create an 'images' folder in your Space and upload some examples.
+example_list = [
+    ["images/example_apple.jpg"],      # <-- REPLACE with actual paths in your Space repo
+    ["images/example_banana.jpg"],     # <-- REPLACE with actual paths in your Space repo
+    ["images/example_strawberry.jpg"]  # <-- REPLACE with actual paths in your Space repo
+]
+# Check if example files exist, otherwise provide empty list
+if not all(os.path.exists(ex[0]) for ex in example_list):
+     print("Warning: Example image paths not found. Clearing examples.")
+     example_list = []
+
+
+# Define Title, Description, and Article for the Gradio app
+title = "Fruit Classifier 🍎🍌🍓"
+description = """
+Upload an image of a fruit or use one of the examples below.
+This demo uses an EfficientNet-B0 model fine-tuned on the Fruits-360 dataset
+(with merged classes) to predict the fruit type.
+Model hosted on Hugging Face Hub: [{MODEL_REPO_ID}](https://huggingface.co/{MODEL_REPO_ID})
+""".format(MODEL_REPO_ID=MODEL_REPO_ID) # Format description with repo ID
+article = """
+<div style='text-align: center;'>
+Model trained using PyTorch and tracked with Neptune.ai. |
+<a href='https://huggingface.co/{MODEL_REPO_ID}' target='_blank'>Model Repository</a> |
+Built with Gradio
+</div>
+""".format(MODEL_REPO_ID=MODEL_REPO_ID)
+
+# Create and launch the interface
+if model is not None: # Only launch if model loaded successfully
+    iface = gr.Interface(
+        fn=predict,
+        inputs=gr.Image(type="pil", label="Upload Fruit Image"),
+        outputs=gr.Label(num_top_classes=5, label="Predictions"), # Show top 5 predictions
+        title=title,
+        description=description,
+        article=article,
+        examples=example_list,
+        allow_flagging="never" # Optional: disable flagging
+    )
+
+    iface.launch()
+else:
+    print("Gradio interface not launched due to model loading failure.")
+