cnn-classifier-app added

app.py

@@ -0,0 +1,30 @@
+import gradio as gr
+from core.predict import ImageClassifier
+import os
+from PIL import Image
+
+cwd = os.getcwd()
+model_path = os.path.join(cwd,'model','cnn-trained-model.pth')
+class_name = {0 : 'Cat' ,1 : 'Dog' ,2 : 'person'}
+classifier = ImageClassifier(model_path=model_path,class_name=None)
+
+
+def classify_image(image):
+    image_path = 'uploaded_image.jpg'
+    image.save(image_path)
+
+    label,output_path = classifier.predict(image_path)
+
+    return label,Image.open(output_path)
+
+demo = gr.Interface(
+    fn=classify_image,
+    inputs = gr.Image(type='pil'),
+    outputs=[gr.Textbox(label="Prediction"),gr.Image(label="Labeled Image")],
+    title="Image Classification Gradio app",
+    description="Upload an Image to classify it as Dog,Cat or Person"
+)
+
+
+if __name__ == '__main__':
+    demo.launch()

core/predict.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from PIL import Image
+import cv2 as cv
+import os
+
+
+
+class CustomCnnModel(nn.Module):
+    def __init__(self,input_dim,num_classes):
+        super(CustomCnnModel,self).__init__()
+        self.input_dim = input_dim
+        self.num_classes = num_classes
+
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(in_channels=3,out_channels=32,kernel_size=3,stride=1,padding=1), # 122x128x3 --> 3x3x3x32 --> wxhx32
+            nn.BatchNorm2d(num_features=32),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2,stride=2),
+
+            nn.Conv2d(in_channels=32,out_channels=64,kernel_size=3,stride=1,padding=1),
+            nn.BatchNorm2d(num_features=64),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2,stride=2),
+
+            nn.Conv2d(in_channels=64,out_channels=128,kernel_size=3,stride=1,padding=1),
+            nn.BatchNorm2d(num_features=128),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2,stride=2),
+
+            nn.Conv2d(in_channels=128,out_channels=256,kernel_size=3,stride=1,padding=1),
+            nn.BatchNorm2d(num_features=256),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2,stride=2),
+        )
+
+        self._to_linear = None
+        self._get_conv_output(self.input_dim)
+
+
+        self.fc_layers = nn.Sequential(
+            nn.Linear(self._to_linear,512),
+            nn.ReLU(),
+            nn.Linear(512,128),
+            nn.ReLU(),
+            nn.Linear(128,self.num_classes)
+        )
+
+
+    def _get_conv_output(self,input_dim=128):
+        with torch.no_grad():
+            dummy_input = torch.zeros(1,3,input_dim,input_dim)
+            output = self.conv_layers(dummy_input)
+            self._to_linear = output.view(1,-1).size(1)
+
+
+
+
+    def forward(self,x):
+        x = self.conv_layers(x)
+        x = x.view(x.size(0),-1)
+        x = self.fc_layers(x)
+        return x
+
+
+
+class ImageClassifier:
+    def __init__(self,model_path,class_name=None):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model = CustomCnnModel(input_dim=128,num_classes=3).to(self.device)
+        self.model.load_state_dict(torch.load(model_path,map_location=self.device))
+        self.model.eval()
+        if class_name is None:
+            self.class_name = {0: 'Cat', 1: 'Dog', 2: 'person'}
+        else:
+            self.class_name = class_name
+
+        self.transform = transforms.Compose([
+        transforms.Resize((128,128)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5]),
+         ]
+        )
+
+    def predict(self,image_path):
+        image = Image.open(image_path).convert('RGB')
+        image_tensor = self.transform(image).unsqueeze(0).to(self.device)
+        
+        with torch.no_grad():
+            output = self.model(image_tensor)
+            _,predicted = torch.max(output,1)
+        
+        label = self.class_name[predicted.item()]
+
+        img = cv.imread(image_path)
+        cv.putText(img,label,(10,30),cv.FONT_HERSHEY_SIMPLEX,1,(255,0,0),2)
+        output_path = 'labeled-image.jpg'
+        cv.imwrite(output_path,img)
+        cwd = os.getcwd()
+        output_path = os.path.join(cwd,output_path)
+        
+        return label,output_path

model/cnn-trained-model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d428f26d1cb36ee748e2ef67df1290a405cc51e31254ae68d4288acb5b7add2d
+size 35393367

requirments.txt

@@ -0,0 +1,5 @@
+gradio
+torch
+torchvision
+opencv-python
+pillow