Update app.py

app.py

@@ -5,21 +5,37 @@ import torchvision
 from PIL import Image
 import gradio as gr
 from huggingface_hub import hf_hub_download
-
-# Configure device
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-transform = torchvision.transforms.Compose([
-    torchvision.transforms.Resize((224, 224)),
-    torchvision.transforms.ToTensor(),
-    torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.225, 0.225, 0.225])
-])
+import torch.nn as nn
+import timm 
 
 REPO_ID = "Raaniel/model-smoke"
 MODEL_FILE_NAME = "model_smoke.pt"
+USE_CUDA = torch.cuda.is_available()
+num_classes = 3
+
+# Download the model
 checkpoint_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE_NAME)
 
-model = torch.jit.load(checkpoint_path)
+# Load the checkpoint
+state = torch.load(checkpoint_path, map_location=torch.device('cuda' if USE_CUDA else 'cpu'))
+
+# Create the model and modify it
+model = timm.create_model('mobilenetv3_small_050', pretrained=True)
+num_features = model.classifier.in_features
+
+# Additional linear and dropout layers
+model.classifier = nn.Sequential(
+    nn.Linear(num_features, 256),  # Additional linear layer
+    nn.ReLU(inplace=True),
+    nn.Dropout(0.5),
+    nn.Linear(256, num_classes)  # Final classification layer
+)
+
+# Load the model weights
+model.load_state_dict(state['weights'])
+
+# Move model to the appropriate device
+device = torch.device('cuda' if USE_CUDA else 'cpu')
 model = model.to(device)
 
 classes = ["chmury", 'inne', "dym"]