Update app.py

app.py

@@ -4,7 +4,8 @@ import os
 import torch
 import torch.nn.functional as F
 import torchvision.transforms as transforms
-import torchvisiona
+import torchvision
+from torchkeras.data import get_url_img
 from torchkeras import plots 
 import numpy as np
 import yaml
@@ -18,14 +19,14 @@ model =  YOLO('Models/haze_detection.pt')
 model = model.to(device)
 
 def load_img (filename):
-    img = Image.open(filename).convert("RGB")
+    if isinstance(img,str):
+        img = get_url_img(img) if img.startswith('http') else Image.open(img).convert('RGB')
     return img
 
 def process_img(image):
-    y = image#.to(device)
 
     with torch.no_grad():
-        result = model(y)
+        result = model.predict(source=image)
 
         if len(result[0].boxes)>0:
             vis = plots.plot_detection(img,boxes=result[0].boxes.boxes,