Update app.py

app.py

@@ -1,55 +1,19 @@
-import torch
-from PIL import Image
-import torchvision.models as models
-from torchvision import transforms
 import gradio as gr
+from fastai.learner import load_learner
+from fastai.vision.core import PILImage
 
-# Load the PyTorch model
-#model = torch.load("model.pth")  # Replace with the path to your trained model
-model = models.resnet18(pretrained=False)  # Instantiate ResNet18 without pre-trained weights
-# Load the saved weights
-model.load_state_dict(torch.load("model.pth"))
-model.eval()  # Set the model to evaluation mode
 
-# Define preprocessing transformations
-preprocess = transforms.Compose([
-    transforms.Resize((224, 224)),  # Resize image to match model's expected input size
-    transforms.ToTensor(),          # Convert image to tensor
-    transforms.Normalize(           # Normalize with ImageNet mean and std
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225]
-    )
-])
+from huggingface_hub import hf_hub_download
 
-# Define the class labels
-class_labels = ["Sailboat", "Fish"]  # Update these labels based on your dataset
 
-# Define the prediction function
-def predict(image):
-    # Preprocess the input image
-    input_tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
+learner = load_learner(hf_hub_download("hmmrg/itisasailboat2","model.pkl"))
 
-    # Perform inference
-    with torch.no_grad():
-        output = model(input_tensor)
-        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+def classify_image(image):
+    img = PILImage.create(image)
+    pred_class, pred_idx, outputs = learner.predict(img)
+    return(f"Predicted class: {pred_class}")
+    
 
-    # Get the predicted class index as an int
-    predicted_index = int(torch.argmax(probabilities).item())
-
-    confidence = probabilities[predicted_index].item()
-
-    return {class_labels[predicted_index]: confidence}
-
-# Create the Gradio interface
-interface = gr.Interface(
-    fn=predict,              # Prediction function
-    inputs=gr.Image(type="pil"),  # Input is an image (handled as a PIL.Image)
-    outputs=gr.Label(num_top_classes=2),  # Output is a label with top 2 class probabilities
-    title="Sailboat or Fish Classifier",
-    description="Upload an image to classify it as either a sailboat or a fish.",
-)
-
-# Launch the app
-if __name__ == "__main__":
-    interface.launch()
+gr.Interface(fn=classify_image,
+             inputs=gr.Image(width=224, height=224),
+             outputs="textbox").launch()