Update app.py

app.py

@@ -1,4 +1,84 @@
 import gradio as gr
+import wandb
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+from torchvision.datasets import ImageFolder
+from torch.utils.data import DataLoader
+import matplotlib.pyplot as plt
+import torchvision
+import matplotlib.pyplot as plt
+import numpy as np
+
+transform = transforms.Compose([
+    transforms.Resize((128, 128)),
+    transforms.ToTensor(),
+    #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomRotation(10),
+    #transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2)
+])
+
+train_dataset = ImageFolder(root='/content/data/train', transform=transform)
+train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+
+val_dataset = ImageFolder(root='/content/data/val', transform=transform)
+val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
+
+test_dataset = ImageFolder(root='/content/data/val', transform=transform)
+test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
+
+class WeatherNet(nn.Module):
+    def __init__(self):
+        super(WeatherNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+        self.fc1 = nn.Linear(128 * 16 * 16, 512)
+        self.fc2 = nn.Linear(512, 11)  # 11 classes
+
+    def forward(self, x):
+        x = self.pool(torch.relu(self.conv1(x)))
+        x = self.pool(torch.relu(self.conv2(x)))
+        x = self.pool(torch.relu(self.conv3(x)))
+        x = x.view(-1, 128 * 16 * 16)
+        x = torch.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+model = WeatherNet()
+
+
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+
+num_epochs = 10
+for epoch in range(num_epochs):
+    model.train()
+    running_loss = 0.0
+    for images, labels in train_loader:
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        running_loss += loss.item()
+    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")
+
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for images, labels in val_loader:
+            outputs = model(images)
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+    print(f'Validation Accuracy: {100 * correct / total:.2f}%')
+
 import wandb
 def predict_image(image):
     image = Image.fromarray(image).convert("RGB")