import gradio as gr
from transformers import pipeline
from PIL import Image
import pytesseract
#import fitz # PyMuPDF
import os
from huggingface_hub import InferenceClient
from phi.agent import Agent, RunResponse
from phi.model.groq import Groq
from pydantic import BaseModel, Field
from typing import Optional
document_types = {
# Wildlife Crime Control Bureau (WCCB)
"CITES Certificate": {
"code": "6261WC",
"description": "A certificate used in the trade of endangered species in accordance with the CITES convention."
},
"Captive Bred/Artificially Propagated Certificate": {
"code": "626WC3",
"description": "A certificate issued by WCCB for species that fall under Appendix 1 of CITES, certifying that the species were bred or propagated artificially."
},
# Central Pollution Control Board (CPCB)
"Extended Producers Responsibility Authorization (EPRA)": {
"code": "651001",
"description": "Authorization issued by CPCB for producers to manage the end-of-life disposal of their products responsibly."
},
# Central Drug Standard Control Organization (CDSCO)
"Invoice": {
"code": "380000",
"description": "A commercial document issued by a seller to a buyer, indicating the products, quantities, and agreed prices for products or services."
},
"Invoice cum Packing List": {
"code": "331000",
"description": "A combined document serving as both an invoice and a packing list, detailing the contents of the shipment."
},
"Master Bill of Lading": {
"code": "704000",
"description": "A bill of lading issued by the carrier or its agent, covering multiple house bills of lading."
},
"House Bill of Lading": {
"code": "714000",
"description": "A bill of lading issued by a freight forwarder or consolidator to individual shippers, covering part of a larger shipment."
},
"Bill of Lading": {
"code": "705000",
"description": "A legal document issued by a carrier to a shipper, detailing the type, quantity, and destination of the goods being carried."
},
"House Waybill": {
"code": "703000",
"description": "A transport document issued by a freight forwarder or consolidator, serving as a receipt for goods shipped."
},
"Tanker Bill of Lading": {
"code": "709000",
"description": "A bill of lading specifically for liquid bulk cargo transported in tankers."
},
"Sea Waybill (Non Negotiable)": {
"code": "710000",
"description": "A non-negotiable transport document evidencing the contract for the carriage of goods by sea."
},
"Inland Waterway Bill of Lading": {
"code": "711000",
"description": "A negotiable transport document for goods transported via inland waterways."
},
"Air Waybill": {
"code": "740000",
"description": "A transport document issued by an airline for the carriage of goods by air."
},
"Master Air Waybill": {
"code": "741000",
"description": "A transport document issued by an airline or consolidator for consolidated shipments from multiple shippers."
},
"Waybill (Non Negotiable)": {
"code": "700000",
"description": "A non-negotiable document evidencing the contract for the transport of cargo."
},
# Other Documents
"CC for Tea Import": {
"code": "911TB1",
"description": "Certificate for tea import issued by the Tea Board."
},
"Kimberley Process Certificate": {
"code": "002KP0",
"description": "A certificate issued by GJEPC for the export and import of rough diamonds, ensuring they are conflict-free."
},
"Import Certificate for Narcotic Drugs": {
"code": "911BN2",
"description": "A certificate issued by CBN for the import of narcotic drugs, ensuring compliance with regulations."
},
"Registration Certificate (Drugs)": {
"code": "101DC1",
"description": "A certificate issued by CDSCO for the registration of drug manufacturers."
},
"Registration Certificate (Cosmetics)": {
"code": "101DC2",
"description": "A certificate issued by CDSCO for the registration of cosmetic products."
},
"Import Permit for Non-Insecticidal Use": {
"code": "911CI3",
"description": "A permit issued by CIB&RC for the import of products intended for non-insecticidal use."
},
"Permit for Import of Livestock Products": {
"code": "911DF1",
"description": "A permit issued by DADF for the import of livestock products, ensuring compliance with health and safety standards."
},
"NOC for Import/Procurement of Aircraft": {
"code": "911DA3",
"description": "A No Objection Certificate issued by DGCA for the import or procurement of aircraft for air transport services."
},
"Licence Form LE-8 (Licence to Import Explosives)": {
"code": "911PE1",
"description": "A license issued by PESO for the import of explosives, ensuring compliance with safety regulations."
},
"License to Import Ammonium Nitrate": {
"code": "911PE2",
"description": "A license issued by PESO for the import of ammonium nitrate, ensuring compliance with safety regulations."
},
"Import Permit": {
"code": "911PQ1",
"description": "A permit issued by PQIS for the import of specific goods, ensuring compliance with regulatory requirements."
},
"No Objection Certificate (NOC) for Natural Rubber": {
"code": "911RB1",
"description": "A No Objection Certificate issued by the Rubber Board for the import of natural rubber."
},
"Registration-Cum-Membership Certificate (RCMC)": {
"code": "101TO2",
"description": "A certificate issued by the Tobacco Board for the registration and membership of exporters."
},
"Certificate of Registration (Publisher/Owner)": {
"code": "101RN1",
"description": "A certificate issued by RNI for the registration of publishers or owners of publications."
},
"No Objection Certificate for Procurement of Radiation Source": {
"code": "911AB1",
"description": "A No Objection Certificate issued by AERB for the procurement of radiation sources."
},
"Import License for Drugs": {
"code": "9111DC",
"description": "A license issued by CDSCO for the import of drugs, ensuring compliance with regulatory requirements."
},
"Certificate of Analysis": {
"code": "001000",
"description": "A certificate providing the results of an analysis of a product, attesting to its quality or purity."
},
"Certificate of Conformity": {
"code": "002000",
"description": "A certificate certifying that a product conforms to predefined specifications or standards."
},
"Certificate of Quality": {
"code": "003000",
"description": "A certificate certifying the quality of goods, services, or processes."
},
"Test Report": {
"code": "004000",
"description": "A report providing the results of tests conducted on a product or material."
},
"Price/Sales Catalogue": {
"code": "005000",
"description": "A document providing pricing and catalogue details for goods and services offered by a seller."
},
"Product Specification Report": {
"code": "006000",
"description": "A report providing detailed specifications of a product."
},
"Federal Label Approval": {
"code": "011000",
"description": "A document relating to federal label approval requirements for packaged products."
},
"Vaccination Certificate": {
"code": "0380AQ",
"description": "An official document proving immunization against certain diseases."
},
"Safety and Hazard Data Sheet": {
"code": "0530HZ",
"description": "A document providing information about the hazards of a material and advice on its safe handling and storage."
},
"Registration Document": {
"code": "101000",
"description": "An official document providing registration details for a business, organization, or individual."
},
"Importer Exporter Certificate": {
"code": "101001",
"description": "A certificate issued by DGFT to authorize an entity to engage in import or export activities."
},
"Purchase Order": {
"code": "105000",
"description": "A document issued by a buyer to a seller, specifying the details of a purchase transaction."
},
"Bond": {
"code": "165000",
"description": "A document indicating a financial guarantee for the payment of duties or performance of obligations."
},
"Crew List Declaration": {
"code": "250000",
"description": "A declaration listing the crew members aboard a conveyance, such as a ship or aircraft."
},
"Fumigation Certificate": {
"code": "2670PQ",
"description": "A certificate attesting that fumigation has been performed on goods to eliminate pests."
},
"Packing List": {
"code": "271000",
"description": "A document specifying the distribution of goods in individual packages within a shipment."
},
"Inspection Report": {
"code": "293000",
"description": "A report informing a party of the results of an inspection conducted on goods or facilities."
},
"Dangerous Goods List": {
"code": "298000",
"description": "A listing of all details of dangerous goods carried, including their classification and handling requirements."
},
"Contract": {
"code": "315000",
"description": "A document evidencing an agreement between parties for the supply of goods or services."
},
"Proforma Invoice": {
"code": "325000",
"description": "A preliminary invoice containing similar information to the final invoice but not claiming immediate payment."
},
"Banker's Guarantee": {
"code": "430000",
"description": "A document in which a bank guarantees the payment of a specified amount to a designated party under certain conditions."
},
"Documentary Credit": {
"code": "465000",
"description": "A document issued by a bank stating that it will pay, accept, or negotiate drafts upon compliance with specified terms and presentation of required documents."
},
"Free Sale Certificate": {
"code": "627000",
"description": "A certificate confirming that a specified product is freely available for sale in the country of origin."
},
"Veterinary Quarantine Certificate": {
"code": "629000",
"description": "A certification that livestock or animal products are kept under health supervision for a specified period as per veterinary quarantine instructions."
},
"Health Certificate": {
"code": "636000",
"description": "A document legalized by a competent authority, declaring that a product is free from pathogens and fit for human consumption."
},
"Food Grade Certificate": {
"code": "637000",
"description": "A document certifying that a product is suitable for use in the food industry."
},
"Certificate of Compliance with OIE Standards": {
"code": "648000",
"description": "A certification that products have been treated in accordance with standards set by the World Organization for Animal Health (OIE)."
},
"Production Facility License": {
"code": "649000",
"description": "A license granted by a competent authority to a production facility for manufacturing specific products."
},
"Consent to Operate/Registration": {
"code": "6490HZ",
"description": "A certificate issued by a State Pollution Control Board to operate a facility, ensuring compliance with environmental regulations."
},
"Certified Copy of the Licensed Premises": {
"code": "6490DC",
"description": "A certified copy of the premises where drugs are stored, issued by a competent authority."
},
"Manufacturing License": {
"code": "651000",
"description": "A license granted by a competent authority to a manufacturer for the production of specific products."
},
"Industrial License": {
"code": "651001",
"description": "A license granted by a competent authority to establish and operate a specific industry."
},
"Phytosanitary Re-export Certificate": {
"code": "6570PQ",
"description": "A certificate attesting that a consignment meets phytosanitary import requirements for re-export."
},
"Rail Consignment Note": {
"code": "720000",
"description": "A transport document constituting a contract for the carriage of goods by rail."
},
"Road Consignment Note": {
"code": "730000",
"description": "A transport document evidencing a contract for the carriage of goods by road."
},
"Container Manifest": {
"code": "788000",
"description": "A document specifying the contents of particular freight containers or other transport units."
},
"Maritime Declaration of Health": {
"code": "797000",
"description": "A document certifying the health condition on board a vessel, valid until a specified date."
},
"Certificate of Registry": {
"code": "798000",
"description": "An official certificate stating the vessel's registry details."
},
"Ship's Stores Declaration": {
"code": "799000",
"description": "A declaration to Customs regarding the contents of the ship's stores, including goods for passengers and crew."
},
"Export Licence": {
"code": "811000",
"description": "A permit issued by a government authority authorizing the exportation of a specified commodity subject to certain conditions."
},
"Phytosanitary Certificate": {
"code": "8510FS",
"description": "A certificate attesting that a consignment meets phytosanitary import requirements."
},
"Sanitary Certificate": {
"code": "852000",
"description": "A document issued by the competent authority in the exporting country, certifying that alimentary and animal products are fit for human consumption."
},
"Veterinary Certificate": {
"code": "853000",
"description": "A document issued by the competent authority in the exporting country, certifying that live animals or birds are free from disease."
},
"Inspection Certificate": {
"code": "856000",
"description": "A document issued by a competent body, certifying that goods have been inspected in accordance with national or international standards."
},
"Pre-Shipment Inspection Certificate": {
"code": "856001",
"description": "A certificate issued prior to shipment, certifying that goods have been inspected and meet specified standards."
},
"Certificate of Origin": {
"code": "861000",
"description": "A document certifying that goods originate from a specific country."
},
"Preference Certificate of Origin": {
"code": "864000",
"description": "A certificate of origin meeting the requirements for preferential treatment under trade agreements."
},
"Certificate of Origin Form GSP": {
"code": "865000",
"description": "A specific form of certificate of origin for goods qualifying for preferential treatment under the Generalized System of Preferences (GSP)."
},
"Consular Invoice": {
"code": "870000",
"description": "A document prepared by an exporter and endorsed by the diplomatic representation of the importing country, used for customs purposes."
},
"Dangerous Goods Declaration": {
"code": "890000",
"description": "A document describing hazardous goods for transport purposes, certifying compliance with safety regulations."
},
"Import Licence": {
"code": "911000",
"description": "A document issued by the competent authority authorizing the import of specific goods under specified conditions."
},
"Value Declaration (GATT Valuation Declaration)": {
"code": "934000",
"description": "A document in which the importer declares the value of goods for customs purposes, including invoice price, freight, insurance, and other costs."
},
"Customs Invoice": {
"code": "935000",
"description": "A document required by customs in the importing country, detailing the invoice price, freight, insurance, and other costs for determining customs value."
},
"ATA Carnet": {
"code": "955000",
"description": "An international customs document used for the temporary admission of goods, serving as a passport for goods."
},
"Rejection Certificate": {
"code": "XXX0FS",
"description": "A certificate issued by the importing country when an export consignment is rejected, providing details of the rejection."
},
"Narcotics Certificate": {
"code": "XXX001",
"description": "A certificate issued by the country of origin certifying that an item does not contain any narcotic or psychotropic substances."
},
"Radio Activity Report": {
"code": "XXX003",
"description": "A report providing the levels of radioactive elements in food items, ensuring compliance with safety standards."
},
"Protocols Approval Letter": {
"code": "XXX000",
"description": "An approval letter containing details of protocols to be followed in the handling of a product, ensuring compliance with regulatory requirements."
}
}
# Load zero-shot classifier pipeline
classifier = pipeline("zero-shot-classification", model="knowledgator/comprehend_it-base")
# Document type labels and descriptions
# document_types = {
# # Wildlife Crime Control Bureau (WCCB)
# "CITES Certificate": {
# "code": "6261WC",
# "description": "A certificate used in the trade of endangered species in accordance with the CITES convention."
# },
# "Captive Bred/Artificially Propagated Certificate": {
# "code": "626WC3",
# "description": "A certificate issued by WCCB for species that fall under Appendix 1 of CITES, certifying that the species were bred or propagated artificially."
# },
# # Central Pollution Control Board (CPCB)
# "Extended Producers Responsibility Authorization (EPRA)": {
# "code": "651001",
# "description": "Authorization issued by CPCB for producers to manage the end-of-life disposal of their products responsibly."
# },
# }
LABELS = list(document_types.keys())
label_info = "\n".join(f"- {label} ({document_types[label]['code']})" for label in LABELS)
print(label_info)
classifier = pipeline("zero-shot-classification", model="knowledgator/comprehend_it-base")
client = InferenceClient("huggingchat/pdf-to-markdown")
import os
os.environ["GROQ_API_KEY"] = os.environ.get("GROQ_API_KEY")
phi_agent = Agent(
name="doc_classifier",
role=f'Classify the OCR extracted text into top 3 matching document types among the given document types only with confidence scores. Document types are {label_info}. Dont forget to mention the document codes also',
instructions=[
"Read OCR extracted document text and classify it into the top 3 document types provided.",
f"Don't choose any document type of your own. Choose only from the given document types: {label_info}"
],
model=Groq(id="deepseek-r1-distill-llama-70b"),
response_model=StructuredOutput,
markdown=True,
debug_mode=True,
show_tool_calls=True,
monitoring=True
)
class StructuredOutput(BaseModel):
doc_type_1: Optional[str] = Field(None)
confidence_1: Optional[float] = Field(None)
doc_type_2: Optional[str] = Field(None)
confidence_2: Optional[float] = Field(None)
doc_type_3: Optional[str] = Field(None)
confidence_3: Optional[float] = Field(None)
def extract_text_from_pdf(file_path):
try:
with open(file_path, "rb") as f:
result = client.post("/predict", files={"pdf_file": f})
extracted_text = result[0] if isinstance(result, (list, tuple)) and len(result) > 0 else ""
return extracted_text.strip()
except Exception as e:
return f"Error extracting text from PDF: {e}"
def extract_text_from_image(image):
try:
text = pytesseract.image_to_string(image)
return text
except Exception as e:
return f"Error extracting text from image: {e}"
def classify_with_bart(text, candidate_labels):
result = classifier(text, candidate_labels=candidate_labels, multi_label=True)
top_scores = list(zip(result['labels'][:3], result['scores'][:3]))
return {lbl: round(score * 100, 2) for lbl, score in top_scores}
def classify_with_phi(text, label_info):
prompt = f'Classify the below text of document: {text} Given Document types: {label_info}'
response: RunResponse = phi_agent.run(prompt)
result = response.content
return {
result.doc_type_1: round(result.confidence_1 , 2),
result.doc_type_2: round(result.confidence_2 , 2),
result.doc_type_3: round(result.confidence_3 , 2)
}
def process_file(input_file, classifier_choice):
file_path = input_file.name
file_ext = os.path.splitext(file_path)[-1].lower()
text = ""
if file_ext in [".jpg", ".jpeg", ".png", ".bmp"]:
image = Image.open(file_path)
text = extract_text_from_image(image)
elif input_file.name.endswith(".pdf"):
text = extract_text_from_pdf(file_path)
else:
return "Unsupported file format", {}, {}
if not text.strip():
return "No text extracted.", {}, {}
if classifier_choice == "BART Zero-Shot Classifier":
scores_dict = classify_with_bart(text, LABELS)
else:
scores_dict = classify_with_phi(text, label_info)
return text, scores_dict, scores_dict
with gr.Blocks() as demo:
gr.Markdown("""
# 🧠e-Sanchit Document Type Classification App
Upload a **PDF or image**, extract text using OCR, and classify it!
""")
with gr.Row():
input_file = gr.File(label="Upload Image or PDF")
classifier_choice = gr.Radio(["BART Zero-Shot Classifier", "Phi Agent Classifier"],value="Phi Agent Classifier",visible=False, label="Choose Classifier")
output_text = gr.Textbox(label="Extracted Text", lines=10)
output_labels = gr.Label(num_top_classes=3, label="Top 3 Predictions with Confidence Scores")
#output_barplot = gr.BarPlot(label="Prediction Scores", x="Document Type", y="Confidence (%)")
classify_button = gr.Button("Run Classification")
classify_button.click(
process_file,
inputs=[input_file, classifier_choice],
outputs=[output_text, output_labels]
)
gr.HTML(f"""
📄 Available Document Types: Total 89 types
""")
demo.launch(debug=True)
# Initialize HuggingFace Inference client for PDF OC