PushkarA07 commited on
Commit
756262f
1 Parent(s): 001297b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -1,5 +1,6 @@
1
  from fastai.vision.models.unet import DynamicUnet
2
  from torchvision.models.resnet import resnet18
 
3
  from fastai.vision.learner import create_body
4
  import streamlit as st
5
  from PIL import Image
@@ -355,7 +356,7 @@ from fastai.vision.models.unet import DynamicUnet
355
 
356
  def build_res_unet(n_input=1, n_output=2, size=256):
357
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
358
- body = create_body(resnet18, pretrained=True, n_in=n_input, cut=-2)
359
  net_G = DynamicUnet(body, n_output, (size, size)).to(device)
360
  return net_G
361
 
 
1
  from fastai.vision.models.unet import DynamicUnet
2
  from torchvision.models.resnet import resnet18
3
+ from fastai.vision.models import resnet18
4
  from fastai.vision.learner import create_body
5
  import streamlit as st
6
  from PIL import Image
 
356
 
357
  def build_res_unet(n_input=1, n_output=2, size=256):
358
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
359
+ body = create_body(resnet18(pretrained=True), n_in=n_input, cut=-2)
360
  net_G = DynamicUnet(body, n_output, (size, size)).to(device)
361
  return net_G
362