xelpmocAI/documentClassification

1import re2import gradio as gr3 4import torch5 6 7from transformers import AutoFeatureExtractor, AutoModelForImageClassification8 9extractor = AutoFeatureExtractor.from_pretrained("DunnBC22/dit-base-Business_Documents_Classified_v2")10 11model = AutoModelForImageClassification.from_pretrained("DunnBC22/dit-base-Business_Documents_Classified_v2")12 13device = "cuda" if torch.cuda.is_available() else "cpu"14model.to(device)15 16 17def classify_documents(image):18    # input_image = image.convert("RGB")19    inputs = extractor(images=image, return_tensor='pt')20    tensors = torch.from_numpy(inputs.pixel_values[0]).unsqueeze(0)21    model_output = model(tensors).logits22    max_index = torch.argmax(model_output)23 24    document_class = model.config.id2label[max_index.item()]25 26    return {27        "result" : str(document_class)28    }29 30article = "<p style='text-align: center'><a href='https://www.xelpmoc.in/' target='_blank'>Made by Xelpmoc</a></p>"31 32demo = gr.Interface(33    fn=classify_documents,34    inputs="image",35    outputs="json",36    title="Document Classification",37    article=article,38    enable_queue=True,39    examples=[40 41        ["./test_images/email_image_2.jpg"],42        ["./test_images/form_image_3.jpg"]43    ],44    cache_examples=False)45 46demo.launch()