updated app.py

app.py

@@ -1,43 +1,41 @@
-import gradio as gr
-import torch
-from torchvision import transforms
-
-from model import EfficientNet
-# Load the PyTorch model
-# model = torch.load('efficientnet_b2-epoch08-val_loss0_02_cat-vs-dog_clasifier.pt')
-model = EfficientNet.load_from_checkpoint("./efficientnet_b2-epoch49-val_loss0.02.ckpt")
-model.eval()
-# Define the image preprocessing
-transform = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-])
-
-# Define the prediction function
-def predict(image):
-    image = transform(image).unsqueeze(0)
-    with torch.no_grad():
-        output = model(image)
-        print(output)
-        probabilities = torch.softmax(output, dim=1)
-        print(probabilities)
-        _, predicted = torch.max(output.data, 1)
-        prediction = 'Cat' if predicted == 0 else 'Dog'
-        confidence = probabilities.max().item()
-        result = {prediction: confidence}
-    return result
-
-# Create the Gradio interface
-app = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=1),
-    title="Dog vs Cat Classifier",
-    description="Upload an image to classify whether it's a dog or a cat.",
-    allow_flagging='never'
-)
-
-# Launch the app
+import gradio as gr
+import torch
+from torchvision import transforms
+
+from model import EfficientNet
+# Load the PyTorch model
+# model = torch.load('efficientnet_b2-epoch08-val_loss0_02_cat-vs-dog_clasifier.pt')
+model = EfficientNet.load_from_checkpoint("./efficientnet_b2-epoch49-val_loss0.02.ckpt")
+model.eval()
+# Define the image preprocessing
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+# Define the prediction function
+def predict(image):
+    image = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output = model(image)
+        probabilities = torch.softmax(output, dim=1)
+        _, predicted = torch.max(output.data, 1)
+        prediction = 'Cat' if predicted == 0 else 'Dog'
+        confidence = probabilities.max().item()
+        result = {prediction: confidence}
+    return result
+
+# Create the Gradio interface
+app = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=1),
+    title="Dog vs Cat Classifier",
+    description="Upload an image to classify whether it's a dog or a cat.",
+    allow_flagging='never'
+)
+
+# Launch the app
 app.launch()