Update app.py

app.py

@@ -7,7 +7,7 @@ import cv2
 from PIL import Image
 from model import create_model
 model,transform=create_model(num_of_classes=3)
-model.load_state_dict(torch.load("fire_smoke_weight.pth", map_location=torch.device("cpu")))
+model.load_state_dict(torch.load("fire_smoke_weights.pth", map_location=torch.device("cpu")))
 model.eval()
 def classify_video(video):
     cap = cv2.VideoCapture(video)
@@ -27,7 +27,7 @@ def classify_video(video):
             pred = output.argmax().item()
             predictions.append(pred)
     cap.release()
-    class_names=['DEFAULT', 'FIRE', 'SMOKE']
+    class_names=['DEFAULT', 'FIRE Spotted', 'SMOKE Spotted']
     for i in predictions:
         if i == 1:
             Fire.append(i)
@@ -35,14 +35,10 @@ def classify_video(video):
             Smoke.append(i)
         else:
             Default.append(i)
-    if len(Fire)>10 and len(Smoke)>8:
-        return f"Spotted {class_names[1]} and {class_names[2]}"
-    elif len(Fire)>10:
-        return f"Spotted {class_names[1]}"
-    elif len(Smoke)>8:
-        return f"Spotted {class_names[2]}"
+    if len(Fire)>5 and len(Smoke)>5:
+        return f"Fire and Smoke"
     else:
-        return f"Spotted {class_names[0]}"
+        return f"Spotted {class_names[max(predictions)]}"
 Description="An MobileNET model trained to classify Fire and Smoke through Videos"
 Article="Created at jupyter NoteBook with GPU NVIDIA_GeForce_MX350"
 gr.Interface(