added pipeline

app.py

@@ -2,15 +2,45 @@ import requests
 import gradio as gr
 import torch
 from timm.data import resolve_data_config
-from timm.data.transforms_factory import create_transform 
+from torchvision.models import alexnet
+import torch.nn as nn
+from torchvision import transforms
 
 LABELS = {0:'Cat', 1:'Dog'}
-model = torch.load('CatVsDogsModel.pth',map_location='cpu')
 
-transform = create_transform(**resolve_data_config({},model=model))
+
+model = alexnet(pretrained=True)
+for param in model.parameters():
+    param.requires_grad = False
+
+# Add a avgpool here
+avgpool = nn.AdaptiveAvgPool2d((7, 7))
+
+# Replace the classifier layer
+# to customise it according to our output
+model.classifier = nn.Sequential(
+    nn.Linear(256 * 7 * 7, 1024),
+    nn.Linear(1024, 256),
+    nn.Linear(256, 2),
+)
+
+checkpoint = torch.load(
+    "CatVsDogsModel.pth", map_location=torch.device("cpu")
+)
+
+model.load_state_dict(checkpoint["state_dict"])
+model = model.to('cpu')
+
+
+
+transform = transforms.Compose(
+    [transforms.Resize((128, 128)), transforms.ToTensor()]
+)
+
 
 
 def predict(img):
+    img = transform(img).to('cpu')
     img = img.convert('RGB')
     img = transform(img).unsqueeze(0)
     with torch.no_grad():