Thomas.Chaigneau commited on
Commit
a6a9972
1 Parent(s): 9a5a46c

update description

Browse files
Files changed (1) hide show
  1. main.py +11 -5
main.py CHANGED
@@ -6,8 +6,7 @@ from model import DocuGAN
6
 
7
 
8
  chk_path = "best_model.ckpt"
9
- model = DocuGAN(hidden_size=64, num_channel=1, latent_size=128, batch_size=128)
10
- # model = DocuGAN.load_from_checkpoint(chk_path, strict=False)
11
  model.eval()
12
  transform = T.ToPILImage()
13
 
@@ -17,6 +16,7 @@ def fn(seed: int = 42):
17
  noise = torch.randn(1, 128, 1, 1)
18
  with torch.no_grad():
19
  pred = model(noise)
 
20
  img = transform(pred.squeeze(1))
21
  return img
22
 
@@ -29,8 +29,14 @@ gr.Interface(
29
  outputs='image',
30
  examples=[],
31
  enable_queue=True,
32
- title="DocuGAN",
33
- description="Select random seed and click on Submit to generate a new Document",
34
- article="DocuGAN, Document Generator by ChainYo",
 
 
 
 
 
 
35
  css=".panel { padding: 5px } .moflo-link { color: #999 }"
36
  ).launch()
 
6
 
7
 
8
  chk_path = "best_model.ckpt"
9
+ model = DocuGAN.load_from_checkpoint(chk_path, strict=False)
 
10
  model.eval()
11
  transform = T.ToPILImage()
12
 
 
16
  noise = torch.randn(1, 128, 1, 1)
17
  with torch.no_grad():
18
  pred = model(noise)
19
+ pred = pred.mul(0.5).add(0.5)
20
  img = transform(pred.squeeze(1))
21
  return img
22
 
 
29
  outputs='image',
30
  examples=[],
31
  enable_queue=True,
32
+ title="📄 DocuGAN - This document doesn't exist",
33
+ description="Select your seed and click on `Submit` to generate a new document",
34
+ article="""
35
+ The SN-GAN model has been trained on the `invoice` part of RVL-CDIP dataset, available [here](https://huggingface.co/datasets/ChainYo/rvl-cdip-invoice).
36
+
37
+ You can see the full implementation on the dedicated [Colab notebook](https://colab.research.google.com/drive/1u6Ct3KnNl7rcgla0268cp-XGTMmVUuJL?usp=sharing).
38
+
39
+ Made with ❤️ by [@ChainYo](-https://huggingface.co/ChainYo)
40
+ """,
41
  css=".panel { padding: 5px } .moflo-link { color: #999 }"
42
  ).launch()