Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from PIL import Image | |
| import pdf2image | |
| from transformers import ( | |
| LayoutLMv3FeatureExtractor, | |
| LayoutLMv3TokenizerFast, | |
| LayoutLMv3Processor, | |
| LayoutLMv3ForSequenceClassification | |
| ) | |
| import torch | |
| import logging | |
| import traceback | |
| import os | |
| import tempfile | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| class DocumentClassifier: | |
| """Class to classify document types based on layout and text content""" | |
| def __init__(self): | |
| try: | |
| self.feature_extractor = LayoutLMv3FeatureExtractor() | |
| self.tokenizer = LayoutLMv3TokenizerFast.from_pretrained( | |
| "microsoft/layoutlmv3-base" | |
| ) | |
| self.processor = LayoutLMv3Processor(self.feature_extractor, self.tokenizer) | |
| self.model = LayoutLMv3ForSequenceClassification.from_pretrained( | |
| "RavindiG/layoutlmv3-document-classification-v2" | |
| ) | |
| self.model.eval() | |
| self.id2label = { | |
| 0: 'Financial Report', | |
| 1: 'Invoice or Receipt', | |
| 2: 'Legal Document', | |
| 3: 'Medical Record', | |
| 4: 'Research Paper' | |
| } | |
| logger.info("DocumentClassifier initialized successfully.") | |
| except Exception as e: | |
| logger.error(f"Error during initialization: {e}") | |
| logger.error(traceback.format_exc()) | |
| raise | |
| def predict_document_class(self, image): | |
| try: | |
| logger.info(f"Image size: {image.size}, mode: {image.mode}") | |
| if image.mode != 'RGB': | |
| logger.info(f"Converting image from {image.mode} to RGB") | |
| image = image.convert('RGB') | |
| logger.info("Processing image...") | |
| encoded_inputs = self.processor(image, max_length=512, return_tensors="pt") | |
| device = next(self.model.parameters()).device | |
| for k, v in encoded_inputs.items(): | |
| encoded_inputs[k] = v.to(device) | |
| logger.info("Running model inference...") | |
| with torch.no_grad(): | |
| outputs = self.model(**encoded_inputs) | |
| logits = outputs.logits | |
| logger.info(f"Logits: {logits}") | |
| threshold = getattr(self, "confidence_threshold", 1.0) | |
| max_logit = logits.max().item() | |
| # Get probabilities for all classes | |
| probabilities = torch.nn.functional.softmax(logits, dim=-1)[0] | |
| if max_logit < threshold: | |
| logger.info("Max logit below threshold, returning 'Other'") | |
| predicted_label = "Other" | |
| confidence_dict = {"Other": 1.0} | |
| else: | |
| predicted_class_idx = logits.argmax(-1).item() | |
| predicted_label = self.id2label.get(predicted_class_idx, "Other") | |
| # Create confidence dictionary for all classes | |
| confidence_dict = { | |
| self.id2label[i]: probabilities[i].item() | |
| for i in range(len(self.id2label)) | |
| } | |
| logger.info(f"Predicted label: {predicted_label}") | |
| return predicted_label, confidence_dict | |
| except Exception as e: | |
| logger.error(f"Error in predict_document_class: {e}") | |
| logger.error(traceback.format_exc()) | |
| raise | |
| def classify_document(self, file_path): | |
| """Classify document type based on layout and extracted text""" | |
| try: | |
| logger.info(f"Processing file: {file_path}") | |
| if file_path.lower().endswith('.pdf'): | |
| logger.info("Converting PDF to image...") | |
| images = pdf2image.convert_from_path( | |
| file_path, | |
| dpi=120, | |
| fmt="RGB", | |
| first_page=1, | |
| last_page=1 | |
| ) | |
| if not images: | |
| raise ValueError("No images extracted from PDF") | |
| image = images[0] | |
| logger.info(f"PDF converted successfully, image size: {image.size}") | |
| else: | |
| logger.info("Opening image file...") | |
| image = Image.open(file_path) | |
| logger.info(f"Image opened successfully, size: {image.size}") | |
| doc_type, confidence = self.predict_document_class(image) | |
| logger.info(f"Classification successful: {doc_type}") | |
| return doc_type, confidence, image | |
| except Exception as e: | |
| logger.error(f"Error classifying document {file_path}: {e}") | |
| logger.error(traceback.format_exc()) | |
| return "Error", {"Error": 1.0}, None | |
| # Initialize the classifier | |
| logger.info("Initializing DocumentClassifier...") | |
| classifier = DocumentClassifier() | |
| logger.info("Classifier ready!") | |
| def classify_upload(file): | |
| """Gradio interface function for file upload""" | |
| if file is None: | |
| return "No file uploaded", {}, None | |
| try: | |
| # Get the file path | |
| file_path = file.name | |
| # Classify the document | |
| doc_type, confidence, image = classifier.classify_document(file_path) | |
| return doc_type, confidence, image | |
| except Exception as e: | |
| logger.error(f"Error in classify_upload: {e}") | |
| logger.error(traceback.format_exc()) | |
| return f"Error: {str(e)}", {}, None | |
| # Create Gradio interface | |
| with gr.Blocks(title="Document Classifier") as demo: | |
| gr.Markdown( | |
| """ | |
| # 📄 Document Type Classifier | |
| Upload a document image or PDF to classify it into one of the following categories: | |
| - Financial Report | |
| - Invoice or Receipt | |
| - Legal Document | |
| - Medical Record | |
| - Research Paper | |
| **Note:** Only the first page of PDFs will be analyzed. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| file_input = gr.File( | |
| label="Upload Document (Image or PDF)", | |
| file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"] | |
| ) | |
| classify_btn = gr.Button("Classify Document", variant="primary") | |
| with gr.Column(): | |
| label_output = gr.Textbox(label="Document Type", lines=2) | |
| confidence_output = gr.Label(label="Confidence Scores", num_top_classes=5) | |
| image_output = gr.Image(label="Document Preview", type="pil") | |
| # Examples | |
| gr.Markdown("### Examples") | |
| gr.Markdown("Upload your document using the file uploader above.") | |
| # Connect the button to the function | |
| classify_btn.click( | |
| fn=classify_upload, | |
| inputs=[file_input], | |
| outputs=[label_output, confidence_output, image_output] | |
| ) | |
| # Also trigger on file upload | |
| file_input.change( | |
| fn=classify_upload, | |
| inputs=[file_input], | |
| outputs=[label_output, confidence_output, image_output] | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch() |