Update app.py

app.py

@@ -2,41 +2,96 @@ import gradio as gr
 import torch
 from torchvision import transforms
 from PIL import Image
-import numpy as np
+import torch.nn as nn
+import os
+from torchvision import models
 
-# Load the trained model
+# Custom Residual Block
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+        # Skip connection
+        self.skip = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.skip_bn = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        identity = self.skip(x)
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.bn2(self.conv2(x))
+        x += identity  # Add skip connection
+        x = self.relu(x)
+        return x
+
+# EfficientNet Model with Novelty (Residual Block)
+class EfficientNetWithNovelty(nn.Module):
+    def __init__(self, num_classes):
+        super(EfficientNetWithNovelty, self).__init__()
+
+        # Load pre-trained EfficientNet-B0 model
+        self.model = models.efficientnet_b0(pretrained=True)
+
+        # Modify the final classifier layer for our number of classes
+        self.model.classifier[1] = nn.Linear(self.model.classifier[1].in_features, num_classes)
+
+        # Add the custom residual block after the EfficientNet feature extractor
+        self.residual_block = ResidualBlock(1280, 1280)  # 1280 is the output channels from EfficientNet B0
+
+    def forward(self, x):
+        # Pass through the EfficientNet feature extractor
+        x = self.model.features(x)  # Access feature extraction part
+        
+        # Pass through the custom residual block
+        x = self.residual_block(x)
+        
+        # Flatten the output to feed into the classifier
+        x = x.mean([2, 3])  # Global Average Pooling
+        x = self.model.classifier(x)  # Pass through the final classifier layer
+        
+        return x
+
+# Load the model and weights
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = YourModel()  # Replace 'YourModel' with your actual model class
-model.load_state_dict(torch.load('D:/Dataset/Cricket Bowl  Grip/final_model.pth'))
+
+# Update this path with your model path
+model_path = 'final_model.pth'
+num_classes = 10  # Assuming you have 10 classes, update based on your dataset
+
+model = EfficientNetWithNovelty(num_classes)
+model.load_state_dict(torch.load(model_path, map_location=device))
 model.to(device)
 model.eval()
 
-# Define the transformation to be applied to the input image
+# Define image transformations (same as during training)
 transform = transforms.Compose([
-    transforms.Resize((224, 224)),  # Resize image to fit your model's input size
-    transforms.ToTensor(),  # Convert image to tensor
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # Example normalize values for ImageNet
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Define a function for making predictions
+# Define the prediction function for Gradio
 def predict(image):
     image = Image.fromarray(image)  # Convert numpy array to PIL Image
     image = transform(image).unsqueeze(0)  # Apply transformations and add batch dimension
     image = image.to(device)
-    
+
     with torch.no_grad():
         outputs = model(image)
         _, predicted = torch.max(outputs, 1)
-        
-    # Map predicted label to class name
+
+    # Class names for your classification
     class_names = ['OUTSWING', 'STRAIGHT', 'BACK_OF_HAND', 'CARROM', 'CROSSSEAM', 
                    'GOOGLY', 'INSWING', 'KNUCKLE', 'LEGSPIN', 'OFFSPIN']
     predicted_label = class_names[predicted.item()]
-    
     return predicted_label
 
 # Create the Gradio Interface
-iface = gr.Interface(fn=predict, 
+iface = gr.Interface(fn=predict,
                      inputs=gr.Image(type="numpy"),  # Accepts image input
                      outputs=gr.Text(),  # Output the predicted class label
                      live=True)  # live=True enables prediction while image is being uploaded