paugar commited on
Commit
304ad99
·
verified ·
1 Parent(s): fab93aa

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -4
app.py CHANGED
@@ -4,7 +4,8 @@ import os
4
  import torch
5
  import torch.nn.functional as F
6
  import torchvision.transforms as transforms
7
- import torchvisiona
 
8
  from torchkeras import plots
9
  import numpy as np
10
  import yaml
@@ -18,14 +19,14 @@ model = YOLO('Models/haze_detection.pt')
18
  model = model.to(device)
19
 
20
  def load_img (filename):
21
- img = Image.open(filename).convert("RGB")
 
22
  return img
23
 
24
  def process_img(image):
25
- y = image#.to(device)
26
 
27
  with torch.no_grad():
28
- result = model(y)
29
 
30
  if len(result[0].boxes)>0:
31
  vis = plots.plot_detection(img,boxes=result[0].boxes.boxes,
 
4
  import torch
5
  import torch.nn.functional as F
6
  import torchvision.transforms as transforms
7
+ import torchvision
8
+ from torchkeras.data import get_url_img
9
  from torchkeras import plots
10
  import numpy as np
11
  import yaml
 
19
  model = model.to(device)
20
 
21
  def load_img (filename):
22
+ if isinstance(img,str):
23
+ img = get_url_img(img) if img.startswith('http') else Image.open(img).convert('RGB')
24
  return img
25
 
26
  def process_img(image):
 
27
 
28
  with torch.no_grad():
29
+ result = model.predict(source=image)
30
 
31
  if len(result[0].boxes)>0:
32
  vis = plots.plot_detection(img,boxes=result[0].boxes.boxes,