Add contour model

app.py

@@ -83,13 +83,17 @@ class Generator(nn.Module):
 
         return out
 
-model1 = Generator(3, 1, 3)
-model1.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
-model1.eval()
+anime = Generator(3, 1, 3)
+anime.load_state_dict(torch.load('anime.pth', map_location=torch.device('cpu')))
+anime.eval()
 
-model2 = Generator(3, 1, 3)
-model2.load_state_dict(torch.load('model2.pth', map_location=torch.device('cpu')))
-model2.eval()
+contour = Generator(3, 1, 3)
+contour.load_state_dict(torch.load('contour.pth', map_location=torch.device('cpu')))
+contour.eval()
+
+opensketch = Generator(3, 1, 3)
+opensketch.load_state_dict(torch.load('opensketch.pth', map_location=torch.device('cpu')))
+opensketch.eval()
 
 def predict(input_img, ver):
     input_img = Image.open(input_img)
@@ -100,10 +104,12 @@ def predict(input_img, ver):
 
     drawing = 0
     with torch.no_grad():
-        if ver == 'style 2':
-            drawing = model2(input_img)[0].detach()
+        if ver == 'anime':
+            drawing = anime(input_img)[0].detach()
+        elif ver == 'contour':
+            drawing = contour(input_img)[0].detach()
         else:
-            drawing = model1(input_img)[0].detach()
+            drawing = opensketch(input_img)[0].detach()
 
     drawing = transforms.ToPILImage()(drawing)
     return drawing
@@ -111,11 +117,11 @@ def predict(input_img, ver):
 title="informative-drawings"
 description="Gradio Demo for line drawing generation. "
 # article = "<p style='text-align: center'><a href='TODO' target='_blank'>Project Page</a> | <a href='codelink' target='_blank'>Github</a></p>"
-examples=[['cat.png', 'style 1'], ['bridge.png', 'style 1'], ['lizard.png', 'style 2'],]
+examples=[['cat.png', 'anime'], ['bridge.png', 'contour'], ['lizard.png', 'opensketch'],]
 
 
 iface = gr.Interface(predict, [gr.Image(type='filepath'),
-    gr.Radio(['style 1','style 2'], type="value", value='style 1', label='version')],
+    gr.Radio(['anime','opensketch','contour'], type="value", value='contour', label='version')],
     gr.Image(type="pil"), title=title,description=description,examples=examples)
 
 iface.launch()

contour.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8b6ec6db973dce9e9a455115514be72f569081651c561a1a8d96d78bb3ece5f
+size 17173511