Project1/app.py at main · Protagonist-8/Project1 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import streamlit as st
import torch
from torch import nn

from torchvision import transforms
from PIL import Image

class TinyVGG(nn.Module):
  def __init__(self,
               input_shape:int,
               hidden_units:int,
               output_shape:int):
    super().__init__()
    self.conv_block_1=nn.Sequential(
        nn.Conv2d(in_channels=input_shape,
                  out_channels=hidden_units,
                  kernel_size=3,
                  stride=1,
                  padding=0),
        nn.ReLU(),
        nn.Conv2d(in_channels=hidden_units,
                  out_channels=hidden_units,
                  kernel_size=3,
                  stride=1,
                  padding=0),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2,
                     stride=2)
    )

    self.conv_block_2=nn.Sequential(
      nn.Conv2d(in_channels=hidden_units,
                out_channels=hidden_units,
                kernel_size=3,
                stride=1,
                padding=0),
      nn.ReLU(),
      nn.Conv2d(in_channels=hidden_units,
                out_channels=hidden_units,
                kernel_size=3,
                stride=1,
                padding=0),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size=2,
                    stride=2)
  )

    self.classifier=nn.Sequential(
        nn.Flatten(),
        nn.Linear(in_features=hidden_units*13*13,
                  out_features=output_shape)
    )

  def forward(self,x):
    x=self.conv_block_1(x)
    # print(x.shape)
    x=self.conv_block_2(x)
    # print(x.shape)
    x=self.classifier(x)
    # print(x.shape)
    return x


model=TinyVGG(input_shape=3,
              hidden_units=10,
              output_shape=3)

model.load_state_dict(torch.load('baseline_vgg.pt'))
model.eval()

transform=transforms.Compose([
  transforms.Resize((64,64)),
  transforms.ToTensor()
])
class_dict={0:'BlackSpot', 1:'DowneyMildew', 2:'FreshLeaf'}

st.title("Leaf Disease Detection")
img=st.file_uploader("Upload an Image",type=['jpg','png'])

if img:
  image=Image.open(img)
  st.image(image, caption='Uploaded Image', use_column_width=True)

  img_tensor=transform(image).unsqueeze(0)

  with torch.inference_mode():
    output=model(img_tensor)
  pred_class=output.argmax(dim=1).item()

  label=class_dict[pred_class]
  if(label=='FreshLeaf'):
    st.write("The Leaf is healthy!")
  else:
    st.write(f'Predicted disease: {label}')