first commit - create app.py

app.py

@@ -0,0 +1,73 @@
+import gradio as gr
+import torch
+from torchvision import transforms
+from PIL import Image
+
+class CustomEfficientNet(nn.Module):
+    def __init__(self, num_classes, num_layers, neurons_per_layer):
+        super(CustomEfficientNet, self).__init__()
+        self.base_model = models.efficientnet_b0(pretrained=True)
+        in_features = self.base_model.classifier[1].in_features
+        self.base_model.classifier = nn.Identity()  # Remove the existing classifier
+        
+        # Define custom layers
+        layers = []
+        for _ in range(num_layers):
+            layers.append(nn.Linear(in_features, neurons_per_layer))
+            layers.append(nn.ReLU())
+            layers.append(nn.Dropout(0.5))
+            in_features = neurons_per_layer
+
+        layers.append(nn.Linear(neurons_per_layer, num_classes))
+        self.custom_classifier = nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.base_model(x)
+        x = x.view(x.size(0), -1)  # Flatten the tensor
+        x = self.custom_classifier(x)
+        return x
+
+def create_model(num_classes, num_layers, neurons_per_layer):
+    model = CustomEfficientNet(num_classes, num_layers, neurons_per_layer)
+    return model
+
+def load_model(path, num_classes, num_layers, neurons_per_layer):
+    model = create_model(num_classes, num_layers, neurons_per_layer)
+    model.load_state_dict(torch.load(path, map_location=torch.device('cpu')))
+    model.eval()
+    return model
+
+# Parameters
+num_classes = 52
+num_layers = 3
+neurons_per_layer = 1024
+
+# Load the model
+model = load_model('card_classification_model.pth', num_classes, num_layers, neurons_per_layer)
+
+# Define the transformation
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+# Class names
+class_names = ['class1', 'class2', ..., 'class52']  # Replace with actual class names
+
+def predict(image):
+    image = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        outputs = model(image)
+        _, predicted = torch.max(outputs, 1)
+        return class_names[predicted[0]]
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="label",
+    description="Upload an image to classify"
+)
+
+iface.launch()