nielsr HF staff commited on
Commit
84b9078
1 Parent(s): 248dc7c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -23
app.py CHANGED
@@ -1,29 +1,16 @@
1
- from transformers import AutoFeatureExtractor, AutoModelForImageClassification
2
  import gradio as gr
3
- import torch
4
 
5
- feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
6
- model = AutoModelForImageClassification.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
7
-
8
- def classify_image(image):
9
- encoding = feature_extractor(image, return_tensors="pt")
10
-
11
- with torch.no_grad():
12
- outputs = model(**encoding)
13
- logits = outputs.logits
14
-
15
- predicted_class = model.config.id2label[logits.argmax(-1).item()]
16
-
17
- return predicted_class
18
-
19
- image = gr.inputs.Image(type="pil")
20
- label = gr.outputs.Label(num_top_classes=3)
21
  title = "Document Image Transformer"
22
  description = "Gradio Demo for DiT, the Document Image Transformer pre-trained on IIT-CDIP, a dataset that includes 42 million document images and fine-tuned on RVL-CDIP, a dataset consisting of 400,000 grayscale images in 16 classes, with 25,000 images per class. To use it, simply add your image, or click one of the examples to load them. Read more at the links below."
23
  article = "<p style='text-align: center'><a href='https://huggingface.co/microsoft/dit-base-finetuned-rvlcdip' target='_blank'>Huggingface Model</a></p>"
24
- examples = [
25
- ["coca_cola_advertisement.png"],
26
- ["scientific_publication.png"]
27
- ]
28
 
29
- gr.Interface(fn=classify_image, inputs=image, outputs=label, title=title, description=description, examples=examples, enable_queue=True).launch(debug=True)
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
+ from transformers import pipeline
3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
  title = "Document Image Transformer"
5
  description = "Gradio Demo for DiT, the Document Image Transformer pre-trained on IIT-CDIP, a dataset that includes 42 million document images and fine-tuned on RVL-CDIP, a dataset consisting of 400,000 grayscale images in 16 classes, with 25,000 images per class. To use it, simply add your image, or click one of the examples to load them. Read more at the links below."
6
  article = "<p style='text-align: center'><a href='https://huggingface.co/microsoft/dit-base-finetuned-rvlcdip' target='_blank'>Huggingface Model</a></p>"
 
 
 
 
7
 
8
+ pipe = pipeline(task="image-classification",
9
+ model="microsoft/dit-base-finetuned-rvlcdip")
10
+ gr.Interface.from_pipeline(pipe,
11
+ title=title,
12
+ description=description,
13
+ examples=['coca_cola_advertisement.png', 'scientific_publication.png',],
14
+ article=article,
15
+ enable_queue=True,
16
+ ).launch()