Update app.py

app.py

@@ -53,7 +53,6 @@ class_classes = [
     "End of no passing by vehicles over 3.5 metric tons"
 ]
 
-# 2. model and transfomrs prep
 
 effnetb2, effnetb2_transforms = create_effnetb2_model(43)
 
@@ -65,7 +64,7 @@ effnetb2_transforms_new = torchvision.transforms.Compose([
 
 effnetb2.load_state_dict(torch.load(f="effnetb2_traffic_sign_recognition.pth", map_location=torch.device("cpu")))
 
-# predict function
+
 
 def predict(
     img,
@@ -81,17 +80,17 @@ def predict(
     """
     start = timer()
 
-    # 1. Pre-process
+
     img_t = transform(img).unsqueeze(0)
 
-    # 2. Forward pass
+
     model.eval()
     with torch.inference_mode():
         logits = model(img_t)
-        probs  = torch.softmax(logits, dim=1).squeeze(0)   # shape [43]
+        probs  = torch.softmax(logits, dim=1).squeeze(0)   
 
     # 3. Top-k
-    top_probs, top_idxs = probs.topk(k)                    # tensors of length k
+    top_probs, top_idxs = probs.topk(k)                    
     pred_topk = {
         class_classes[int(idx)]: float(prob)
         for idx, prob in zip(top_idxs, top_probs)
@@ -101,7 +100,7 @@ def predict(
     return pred_topk, pred_time
 
 
-# 4. gradio app
+
 
 import gradio as gr