Update app.py

app.py

@@ -1,29 +1,30 @@
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../drive/MyDrive/Colab Notebooks/cats_vs_dogs.ipynb.
+import torch
+import torchvision
 
-# %% auto 0
-__all__ = ['learn', 'categories', 'image', 'label', 'examples', 'intf', 'classify_image']
+model = torchvision.models.resnet50(pretrained=False)
+model.fc = nn.Linear(model.fc.in_features, num_classes)
+model.load_state_dict(torch.load("model.pth"))
+model.to(device)
+model.eval()
 
-# %% ../drive/MyDrive/Colab Notebooks/cats_vs_dogs.ipynb 10
-from fastai.vision.all import *
 import gradio as gr
-
-def is_cat(x): return x[0].isupper()
-
-
-# %% ../drive/MyDrive/Colab Notebooks/cats_vs_dogs.ipynb 12
-learn = load_learner('pets.pkl')
-
-# %% ../drive/MyDrive/Colab Notebooks/cats_vs_dogs.ipynb 14
-categories = ('Dog', 'Cat')
-
-def classify_image(img):
-    pred,idx,probs = learn.predict(img)
-    return dict(zip(categories, map(float,probs)))
-
-# %% ../drive/MyDrive/Colab Notebooks/cats_vs_dogs.ipynb 15
-image = gr.inputs.Image(shape=(192, 192))
-label = gr.outputs.Label()
-examples = ['dog.jpg']
-
-intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
-intf.launch(inline=False)
+from PIL import Image
+
+# Define the function to make predictions
+def predict(image):
+    image = transform(image).unsqueeze(0).to(device)
+    model.eval()
+    with torch.no_grad():
+        output = model(image)
+        _, predicted = torch.max(output.data, 1)
+        return dataset.classes[predicted.item()]
+
+# Define the input and output components
+image_input = gr.inputs.Image(type="pil", label="Upload Image")
+label_output = gr.outputs.Label()
+
+# Create the interface
+interface = gr.Interface(fn=predict, inputs=image_input, outputs=label_output)
+
+# Launch the interface
+interface.launch()