app.py

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"""
    <hr>
    <h4>📄 Available Document Types: Total 89 types </h4>
    <ul style='line-height:1.6;'>{label_info}</ul>
    """)

    demo.launch(debug=True)

# Initialize HuggingFace Inference client for PDF OC