update app.py

app.py

@@ -13,33 +13,49 @@ class CNNKAN(nn.Module):
     def __init__(self):
         super(CNNKAN, self).__init__()
         self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(32)
         self.pool1 = nn.MaxPool2d(2)
         self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(64)
         self.pool2 = nn.MaxPool2d(2)
-        self.kan1 = KANLinear(64 * 50 * 50, 256)
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.bn3 = nn.BatchNorm2d(128)
+        self.pool3 = nn.MaxPool2d(2)
+        self.dropout = nn.Dropout(0.5)
+        self.kan1 = KANLinear(128 * 25 * 25, 256)
         self.kan2 = KANLinear(256, 1)
 
     def forward(self, x):
-        x = F.selu(self.conv1(x))
+        x = F.selu(self.bn1(self.conv1(x)))
         x = self.pool1(x)
-        x = F.selu(self.conv2(x))
+        x = F.selu(self.bn2(self.conv2(x)))
         x = self.pool2(x)
+        x = F.selu(self.bn3(self.conv3(x)))
+        x = self.pool3(x)
         x = x.view(x.size(0), -1)
+        x = self.dropout(x)
         x = self.kan1(x)
+        x = self.dropout(x)
         x = self.kan2(x)
         return x
 
-# Assuming the model weights are saved in 'model.pth'
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = CNNKAN().to(device)
-model.load_state_dict(torch.load('weights/model_weights_KAN1.pth', map_location=device))
-model.eval()
+def load_model(weights_path, device):
+    model = CNNKAN().to(device)
+    model.load_state_dict(torch.load(weights_path, map_location=device))
+    model.eval()
+    return model
+
+def load_image_from_url(url):
+    response = requests.get(url)
+    img = Image.open(BytesIO(response.content)).convert('RGB')
+    return img
 
-# Define image transformations
-transform = transforms.Compose([
-    transforms.Resize((200, 200)),
-    transforms.ToTensor()
-])
+def preprocess_image(image):
+    transform = transforms.Compose([
+        transforms.Resize((200, 200)),
+        transforms.ToTensor()
+    ])
+    return transform(image).unsqueeze(0)
 
 # Streamlit app
 st.title("Image Classification with CNN-KAN")
@@ -48,10 +64,8 @@ st.sidebar.title("Upload Images")
 uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg", "png", "webp"])
 image_url = st.sidebar.text_input("Or enter image URL...")
 
-def load_image_from_url(url):
-    response = requests.get(url)
-    img = Image.open(BytesIO(response.content)).convert('RGB')
-    return img
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = load_model('weights/best_model_weights_KAN.pth', device)
 
 img = None
 
@@ -66,7 +80,7 @@ elif image_url:
 if img is not None:
     st.image(np.array(img), caption='Uploaded Image.', use_column_width=True)
     if st.button('Predict'):
-        img_tensor = transform(img).unsqueeze(0).to(device)
+        img_tensor = preprocess_image(img).to(device)
 
         with torch.no_grad():
             output = model(img_tensor)

desktop.ini

@@ -0,0 +1,2 @@
+[LocalizedFileNames]
+Screenshot 2024-06-05 110741.png=@Screenshot 2024-06-05 110741.png,0

requirements.txt

@@ -75,4 +75,5 @@ ultralytics==8.1.30
 urllib3==2.2.1
 watchdog==4.0.0
 pafy
-youtube-dl
+youtube-dl
+optuna