Spaces:
Running
Running
PushkarA07
commited on
Commit
•
756262f
1
Parent(s):
001297b
Update app.py
Browse files
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
|
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 |
|