PBJ commited on
Commit
d111ae0
1 Parent(s): c7da235

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -3
app.py CHANGED
@@ -2,8 +2,6 @@ import streamlit as st
2
  from PIL import Image
3
  import cv2 as cv
4
 
5
- # ---------Backend--------------------------------------------------------------
6
-
7
  import os
8
  import glob
9
  import time
@@ -319,7 +317,7 @@ from fastai.vision.models.unet import DynamicUnet
319
 
320
  def build_res_unet(n_input=1, n_output=2, size=256):
321
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
322
- body = create_body(resnet18(), pretrained=True, n_in=n_input, cut=-2)
323
  net_G = DynamicUnet(body, n_output, (size, size)).to(device)
324
  return net_G
325
 
 
2
  from PIL import Image
3
  import cv2 as cv
4
 
 
 
5
  import os
6
  import glob
7
  import time
 
317
 
318
  def build_res_unet(n_input=1, n_output=2, size=256):
319
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
320
+ body = create_body(resnet18, pretrained=True, n_in=n_input, cut=-2)
321
  net_G = DynamicUnet(body, n_output, (size, size)).to(device)
322
  return net_G
323