quantumbit commited on
Commit
819b2d1
Β·
verified Β·
1 Parent(s): 84079a0

Upload 8 files

Browse files
Files changed (8) hide show
  1. Dockerfile +37 -0
  2. app.py +245 -0
  3. config.py +49 -0
  4. inference.py +351 -0
  5. model_manager.py +145 -0
  6. requirements.txt +12 -0
  7. start.sh +26 -0
  8. utils/models/best.pt +3 -0
Dockerfile ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.10-slim
2
+
3
+ WORKDIR /app
4
+
5
+ # Install system dependencies for OpenCV and other libraries
6
+ RUN apt-get update && apt-get install -y \
7
+ git \
8
+ libgl1-mesa-glx \
9
+ libglib2.0-0 \
10
+ libsm6 \
11
+ libxext6 \
12
+ libxrender-dev \
13
+ libgomp1 \
14
+ && rm -rf /var/lib/apt/lists/*
15
+
16
+ # Copy requirements first for better caching
17
+ COPY requirements.txt .
18
+
19
+ # Install Python dependencies
20
+ RUN pip install --no-cache-dir -r requirements.txt
21
+
22
+ # Copy application files
23
+ COPY config.py .
24
+ COPY model_manager.py .
25
+ COPY inference.py .
26
+ COPY app.py .
27
+ COPY utils/ utils/
28
+
29
+ # Expose Hugging Face Spaces default port
30
+ EXPOSE 7860
31
+
32
+ # Health check
33
+ HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
34
+ CMD python -c "import requests; requests.get('http://localhost:7860/health')"
35
+
36
+ # Run the FastAPI application
37
+ CMD ["python", "app.py"]
app.py ADDED
@@ -0,0 +1,245 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ FastAPI Server for Invoice Information Extractor
3
+ Provides REST API for invoice processing
4
+ """
5
+
6
+ from fastapi import FastAPI, File, UploadFile, HTTPException, Form
7
+ from fastapi.responses import JSONResponse
8
+ from fastapi.middleware.cors import CORSMiddleware
9
+ from contextlib import asynccontextmanager
10
+ from typing import Optional
11
+ import tempfile
12
+ import os
13
+ import shutil
14
+
15
+ from config import API_TITLE, API_DESCRIPTION, API_VERSION
16
+ from model_manager import model_manager
17
+ from inference import InferenceProcessor
18
+
19
+
20
+ @asynccontextmanager
21
+ async def lifespan(app: FastAPI):
22
+ """Lifecycle manager - loads models on startup"""
23
+ print("πŸš€ Starting Invoice Information Extractor API...")
24
+ print("=" * 60)
25
+
26
+ # Load models on startup
27
+ try:
28
+ model_manager.load_models()
29
+ print("=" * 60)
30
+ print("βœ… API is ready to accept requests!")
31
+ print("=" * 60)
32
+ except Exception as e:
33
+ print(f"❌ Failed to load models: {str(e)}")
34
+ raise
35
+
36
+ yield
37
+
38
+ # Cleanup on shutdown
39
+ print("πŸ›‘ Shutting down API...")
40
+
41
+
42
+ # Initialize FastAPI app
43
+ app = FastAPI(
44
+ title=API_TITLE,
45
+ description=API_DESCRIPTION,
46
+ version=API_VERSION,
47
+ lifespan=lifespan
48
+ )
49
+
50
+ # Add CORS middleware
51
+ app.add_middleware(
52
+ CORSMiddleware,
53
+ allow_origins=["*"],
54
+ allow_credentials=True,
55
+ allow_methods=["*"],
56
+ allow_headers=["*"],
57
+ )
58
+
59
+
60
+ @app.get("/")
61
+ async def root():
62
+ """Root endpoint - API information"""
63
+ return {
64
+ "name": API_TITLE,
65
+ "version": API_VERSION,
66
+ "status": "running",
67
+ "models_loaded": model_manager.is_loaded(),
68
+ "endpoints": {
69
+ "health": "/health",
70
+ "extract": "/extract (POST)",
71
+ "docs": "/docs"
72
+ }
73
+ }
74
+
75
+
76
+ @app.get("/health")
77
+ async def health_check():
78
+ """Health check endpoint"""
79
+ return {
80
+ "status": "healthy",
81
+ "models_loaded": model_manager.is_loaded()
82
+ }
83
+
84
+
85
+ @app.post("/extract")
86
+ async def extract_invoice(
87
+ file: UploadFile = File(..., description="Invoice image file (JPG, PNG, JPEG)"),
88
+ doc_id: Optional[str] = Form(None, description="Optional document identifier")
89
+ ):
90
+ """
91
+ Extract information from invoice image
92
+
93
+ **Parameters:**
94
+ - **file**: Invoice image file (required)
95
+ - **doc_id**: Optional document identifier (auto-generated from filename if not provided)
96
+
97
+ **Returns:**
98
+ - JSON with extracted fields, confidence scores, and metadata
99
+
100
+ **Example Response:**
101
+ ```json
102
+ {
103
+ "doc_id": "invoice_001",
104
+ "fields": {
105
+ "dealer_name": "ABC Tractors Pvt Ltd",
106
+ "model_name": "Mahindra 575 DI",
107
+ "horse_power": 50,
108
+ "asset_cost": 525000,
109
+ "signature": {"present": true, "bbox": [100, 200, 300, 250]},
110
+ "stamp": {"present": true, "bbox": [400, 500, 500, 550]}
111
+ },
112
+ "confidence": 0.89,
113
+ "processing_time_sec": 3.8,
114
+ "cost_estimate_usd": 0.000528
115
+ }
116
+ ```
117
+ """
118
+
119
+ # Validate file type
120
+ if not file.content_type.startswith("image/"):
121
+ raise HTTPException(
122
+ status_code=400,
123
+ detail="File must be an image (JPG, PNG, JPEG)"
124
+ )
125
+
126
+ # Check if models are loaded
127
+ if not model_manager.is_loaded():
128
+ raise HTTPException(
129
+ status_code=503,
130
+ detail="Models not loaded. Please wait for server initialization."
131
+ )
132
+
133
+ # Save uploaded file to temporary location
134
+ temp_file = None
135
+ try:
136
+ # Create temporary file
137
+ suffix = os.path.splitext(file.filename)[1]
138
+ with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp:
139
+ temp_file = temp.name
140
+ # Write uploaded file content
141
+ shutil.copyfileobj(file.file, temp)
142
+
143
+ # Use filename as doc_id if not provided
144
+ if doc_id is None:
145
+ doc_id = os.path.splitext(file.filename)[0]
146
+
147
+ # Process invoice
148
+ result = InferenceProcessor.process_invoice(temp_file, doc_id)
149
+
150
+ return JSONResponse(content=result)
151
+
152
+ except Exception as e:
153
+ raise HTTPException(
154
+ status_code=500,
155
+ detail=f"Error processing invoice: {str(e)}"
156
+ )
157
+
158
+ finally:
159
+ # Clean up temporary file
160
+ if temp_file and os.path.exists(temp_file):
161
+ try:
162
+ os.unlink(temp_file)
163
+ except:
164
+ pass
165
+
166
+ # Close uploaded file
167
+ file.file.close()
168
+
169
+
170
+ @app.post("/extract_batch")
171
+ async def extract_batch(
172
+ files: list[UploadFile] = File(..., description="Multiple invoice images")
173
+ ):
174
+ """
175
+ Extract information from multiple invoice images
176
+
177
+ **Parameters:**
178
+ - **files**: List of invoice image files
179
+
180
+ **Returns:**
181
+ - JSON array with results for each invoice
182
+ """
183
+
184
+ if not model_manager.is_loaded():
185
+ raise HTTPException(
186
+ status_code=503,
187
+ detail="Models not loaded. Please wait for server initialization."
188
+ )
189
+
190
+ results = []
191
+ temp_files = []
192
+
193
+ try:
194
+ for file in files:
195
+ # Validate file type
196
+ if not file.content_type.startswith("image/"):
197
+ results.append({
198
+ "filename": file.filename,
199
+ "error": "File must be an image"
200
+ })
201
+ continue
202
+
203
+ # Save to temp file
204
+ suffix = os.path.splitext(file.filename)[1]
205
+ with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp:
206
+ temp_file = temp.name
207
+ temp_files.append(temp_file)
208
+ shutil.copyfileobj(file.file, temp)
209
+
210
+ # Process
211
+ try:
212
+ doc_id = os.path.splitext(file.filename)[0]
213
+ result = InferenceProcessor.process_invoice(temp_file, doc_id)
214
+ results.append(result)
215
+ except Exception as e:
216
+ results.append({
217
+ "filename": file.filename,
218
+ "error": str(e)
219
+ })
220
+
221
+ return JSONResponse(content={"results": results})
222
+
223
+ finally:
224
+ # Cleanup
225
+ for temp_file in temp_files:
226
+ if os.path.exists(temp_file):
227
+ try:
228
+ os.unlink(temp_file)
229
+ except:
230
+ pass
231
+
232
+ for file in files:
233
+ file.file.close()
234
+
235
+
236
+ if __name__ == "__main__":
237
+ import uvicorn
238
+
239
+ # Run server
240
+ uvicorn.run(
241
+ "app:app",
242
+ host="0.0.0.0",
243
+ port=7860, # Hugging Face Spaces default port
244
+ reload=False
245
+ )
config.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Configuration settings for Invoice Information Extractor API
3
+ """
4
+
5
+ import os
6
+ from pathlib import Path
7
+
8
+ # Base directories
9
+ BASE_DIR = Path(__file__).resolve().parent
10
+ MODELS_DIR = BASE_DIR / "utils" / "models"
11
+
12
+ # Model paths
13
+ YOLO_MODEL_PATH = MODELS_DIR / "best.pt"
14
+
15
+ # VLM Model Configuration
16
+ VLM_MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
17
+
18
+ # Quantization settings
19
+ QUANTIZATION_CONFIG = {
20
+ "load_in_4bit": True,
21
+ "bnb_4bit_quant_type": "nf4",
22
+ "bnb_4bit_compute_dtype": "float16",
23
+ "bnb_4bit_use_double_quant": True
24
+ }
25
+
26
+ # Image processing settings
27
+ MAX_IMAGE_SIZE = 512 # Maximum dimension for resizing
28
+
29
+ # Detection thresholds
30
+ YOLO_CONFIDENCE_THRESHOLD = 0.25
31
+
32
+ # Validation ranges
33
+ HP_VALID_RANGE = (20, 120)
34
+ ASSET_COST_VALID_RANGE = (100_000, 3_000_000)
35
+
36
+ # Cost calculation
37
+ COST_PER_GPU_HOUR = 0.5 # USD
38
+
39
+ # API settings
40
+ API_TITLE = "Invoice Information Extractor API"
41
+ API_DESCRIPTION = """
42
+ Extract structured information from Indian tractor invoices using AI.
43
+
44
+ **Features:**
45
+ - Extracts dealer name, model name, horse power, and asset cost
46
+ - Detects signatures and stamps with bounding boxes
47
+ - Provides confidence scores and cost estimates
48
+ """
49
+ API_VERSION = "1.0.0"
inference.py ADDED
@@ -0,0 +1,351 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Inference Processor - Handles VLM extraction, validation, and result formatting
3
+ """
4
+
5
+ import torch
6
+ import time
7
+ import json
8
+ import codecs
9
+ import re
10
+ from PIL import Image
11
+ from qwen_vl_utils import process_vision_info
12
+ from typing import Dict, Tuple
13
+
14
+ from config import (
15
+ MAX_IMAGE_SIZE,
16
+ HP_VALID_RANGE,
17
+ ASSET_COST_VALID_RANGE,
18
+ COST_PER_GPU_HOUR
19
+ )
20
+ from model_manager import model_manager
21
+
22
+
23
+ EXTRACTION_PROMPT = """
24
+ You are an expert at reading noisy, handwritten Indian invoices and quotations.
25
+
26
+ Your task is to extract text EXACTLY as it appears in the image.
27
+ Do NOT translate, summarize, normalize, or rewrite any text.
28
+ Preserve the original language (Hindi, Marathi, Kannada, English, etc.).
29
+
30
+ Carefully read the image and extract the following fields.
31
+
32
+ Return ONLY valid JSON in this format:
33
+
34
+ {
35
+ "dealer_name": string,
36
+ "model_name": string,
37
+ "horse_power": number,
38
+ "asset_cost": number
39
+ }
40
+
41
+ Critical rules:
42
+ - Dealer name must be copied exactly from the image in the original language and spelling.
43
+ - Model name must be copied exactly from the image without translation.
44
+ - Do NOT convert regional language text into English.
45
+ - Do NOT expand abbreviations or correct spelling.
46
+ - Only numbers may be normalized.
47
+
48
+ Extraction hints:
49
+ - Asset cost is the total amount, usually the largest number on the page, the total amount after TAX, final price or final cost.
50
+ - Dealer name is usually at the top header or company name.
51
+ - Model name often appears near words like Model, Tractor, Variant.
52
+ - Horse power must come ONLY from explicit HP text, never from model numbers.
53
+ - Horse power may appear as "HP", handwritten like "49 HP", "63hp", "HP-30".
54
+ - Remove commas and currency symbols from numbers only.
55
+ - If handwriting is unclear, make your best reasonable interpretation of the characters β€” but preserve language.
56
+
57
+ Output rules:
58
+ - Output ONLY valid JSON.
59
+ - Do NOT include markdown, explanations, or extra text.
60
+ """
61
+
62
+
63
+ class InferenceProcessor:
64
+ """Handles VLM inference, validation, and result processing"""
65
+
66
+ @staticmethod
67
+ def preprocess_image(image_path: str) -> Image.Image:
68
+ """Load and resize image if needed"""
69
+ image = Image.open(image_path).convert("RGB")
70
+
71
+ # Resize if too large
72
+ if max(image.size) > MAX_IMAGE_SIZE:
73
+ ratio = MAX_IMAGE_SIZE / max(image.size)
74
+ new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
75
+ image = image.resize(new_size, Image.LANCZOS)
76
+ print(f"πŸ”„ Image resized to {new_size}")
77
+
78
+ return image
79
+
80
+ @staticmethod
81
+ def run_vlm_extraction(image: Image.Image) -> Tuple[str, float]:
82
+ """Run VLM model to extract invoice fields"""
83
+ if not model_manager.is_loaded():
84
+ raise RuntimeError("Models not loaded")
85
+
86
+ model = model_manager.vlm_model
87
+ processor = model_manager.processor
88
+
89
+ messages = [
90
+ {
91
+ "role": "user",
92
+ "content": [
93
+ {"type": "image", "image": image},
94
+ {"type": "text", "text": EXTRACTION_PROMPT}
95
+ ]
96
+ }
97
+ ]
98
+
99
+ # Apply chat template
100
+ text = processor.apply_chat_template(
101
+ messages,
102
+ tokenize=False,
103
+ add_generation_prompt=True
104
+ )
105
+
106
+ # Process vision input
107
+ image_inputs, video_inputs = process_vision_info(messages)
108
+ inputs = processor(
109
+ text=[text],
110
+ images=image_inputs,
111
+ videos=video_inputs,
112
+ padding=True,
113
+ return_tensors="pt",
114
+ )
115
+ inputs = inputs.to("cuda")
116
+
117
+ start = time.time()
118
+
119
+ # Generate
120
+ generated_ids = model.generate(**inputs, max_new_tokens=256)
121
+
122
+ latency = time.time() - start
123
+
124
+ # Decode output
125
+ generated_ids_trimmed = [
126
+ out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
127
+ ]
128
+ output_text = processor.batch_decode(
129
+ generated_ids_trimmed,
130
+ skip_special_tokens=True,
131
+ clean_up_tokenization_spaces=False
132
+ )
133
+
134
+ output_text = output_text[0] if isinstance(output_text, list) else output_text
135
+
136
+ # Clean up GPU memory
137
+ del inputs, generated_ids, generated_ids_trimmed
138
+ if torch.cuda.is_available():
139
+ torch.cuda.empty_cache()
140
+
141
+ return output_text, latency
142
+
143
+ @staticmethod
144
+ def extract_json_from_output(text: str) -> Dict:
145
+ """Extract JSON from model output"""
146
+ # Handle single/double backticks
147
+ if text.count('```') in [1, 2]:
148
+ data = text.split('```')[1]
149
+ if data.startswith('json'):
150
+ data = data[4:]
151
+ try:
152
+ return json.loads(codecs.decode(data, "unicode-escape"))
153
+ except:
154
+ pass
155
+
156
+ # Try markdown code blocks
157
+ markdown_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
158
+ if markdown_match:
159
+ try:
160
+ return json.loads(markdown_match.group(1))
161
+ except json.JSONDecodeError:
162
+ pass
163
+
164
+ # Find JSON blocks
165
+ json_matches = re.finditer(r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', text, re.DOTALL)
166
+
167
+ for match in json_matches:
168
+ json_str = match.group(0)
169
+ try:
170
+ parsed = json.loads(json_str)
171
+ # Verify expected keys
172
+ if all(key in parsed for key in ["dealer_name", "model_name", "horse_power", "asset_cost"]):
173
+ return parsed
174
+ except json.JSONDecodeError:
175
+ continue
176
+
177
+ # Fallback
178
+ return {
179
+ "dealer_name": None,
180
+ "model_name": None,
181
+ "horse_power": None,
182
+ "asset_cost": None
183
+ }
184
+
185
+ @staticmethod
186
+ def clean_text(text) -> str:
187
+ """Clean text field"""
188
+ if not text:
189
+ return None
190
+ text = str(text).strip()
191
+ text = re.sub(r"\s+", " ", text)
192
+ return text if len(text) > 1 else None
193
+
194
+ @staticmethod
195
+ def clean_number(num):
196
+ """Clean number field"""
197
+ try:
198
+ if num is None:
199
+ return None
200
+ return int(float(num))
201
+ except:
202
+ return None
203
+
204
+ @staticmethod
205
+ def fix_horse_power(vlm_hp, model_name) -> Tuple:
206
+ """Fix common HP extraction mistakes"""
207
+ # Accept if in valid range
208
+ if vlm_hp is not None and HP_VALID_RANGE[0] <= vlm_hp <= HP_VALID_RANGE[1]:
209
+ return vlm_hp, 1.0
210
+
211
+ # Try extracting from model name
212
+ if model_name:
213
+ match = re.search(r"HP[- ]?(\d+)", model_name, re.I)
214
+ if match:
215
+ hp = int(match.group(1))
216
+ if HP_VALID_RANGE[0] <= hp <= HP_VALID_RANGE[1]:
217
+ return hp, 0.8
218
+
219
+ return None, 0.2
220
+
221
+ @staticmethod
222
+ def validate_asset_cost(cost) -> Tuple:
223
+ """Validate asset cost"""
224
+ if cost is None:
225
+ return None, 0.2
226
+
227
+ cost = InferenceProcessor.clean_number(cost)
228
+
229
+ if ASSET_COST_VALID_RANGE[0] <= cost <= ASSET_COST_VALID_RANGE[1]:
230
+ return cost, 1.0
231
+
232
+ return None, 0.3
233
+
234
+ @staticmethod
235
+ def validate_text_field(text) -> Tuple:
236
+ """Validate text fields"""
237
+ text = InferenceProcessor.clean_text(text)
238
+ if not text or len(text) < 3:
239
+ return None, 0.3
240
+ return text, 1.0
241
+
242
+ @staticmethod
243
+ def validate_prediction(raw_json: Dict) -> Tuple[Dict, float, list]:
244
+ """Validate and fix extracted fields"""
245
+ warnings = []
246
+ confidences = []
247
+
248
+ # Dealer
249
+ dealer, dealer_conf = InferenceProcessor.validate_text_field(raw_json.get("dealer_name"))
250
+ if dealer is None:
251
+ warnings.append("Dealer name invalid")
252
+ confidences.append(dealer_conf)
253
+
254
+ # Model
255
+ model_name, model_conf = InferenceProcessor.validate_text_field(raw_json.get("model_name"))
256
+ if model_name is None:
257
+ warnings.append("Model name invalid")
258
+ confidences.append(model_conf)
259
+
260
+ # Horse Power
261
+ hp_raw = InferenceProcessor.clean_number(raw_json.get("horse_power"))
262
+ hp, hp_conf = InferenceProcessor.fix_horse_power(hp_raw, model_name)
263
+ if hp is None:
264
+ warnings.append("Horse power invalid")
265
+ confidences.append(hp_conf)
266
+
267
+ # Asset Cost
268
+ cost_raw = InferenceProcessor.clean_number(raw_json.get("asset_cost"))
269
+ cost, cost_conf = InferenceProcessor.validate_asset_cost(cost_raw)
270
+ if cost is None:
271
+ warnings.append("Asset cost invalid")
272
+ confidences.append(cost_conf)
273
+
274
+ # Overall field confidence
275
+ field_confidence = round(sum(confidences) / len(confidences), 3)
276
+
277
+ validated = {
278
+ "dealer_name": dealer,
279
+ "model_name": model_name,
280
+ "horse_power": hp,
281
+ "asset_cost": cost
282
+ }
283
+
284
+ return validated, field_confidence, warnings
285
+
286
+ @staticmethod
287
+ def process_invoice(image_path: str, doc_id: str = None) -> Dict:
288
+ """
289
+ Complete invoice processing pipeline
290
+
291
+ Args:
292
+ image_path: Path to invoice image
293
+ doc_id: Document identifier (optional)
294
+
295
+ Returns:
296
+ dict: Complete JSON output with all fields
297
+ """
298
+ total_start = time.time()
299
+
300
+ # Generate doc_id if not provided
301
+ if doc_id is None:
302
+ import os
303
+ doc_id = os.path.splitext(os.path.basename(image_path))[0]
304
+
305
+ # Step 1: Preprocess image
306
+ image = InferenceProcessor.preprocess_image(image_path)
307
+
308
+ # Step 2: YOLO Detection
309
+ signature_info, stamp_info, signature_conf, stamp_conf = model_manager.detect_sign_stamp(image_path)
310
+
311
+ # Step 3: VLM Extraction
312
+ vlm_output, vlm_latency = InferenceProcessor.run_vlm_extraction(image)
313
+
314
+ # Clean up image
315
+ image.close()
316
+ del image
317
+
318
+ # Step 4: Parse JSON
319
+ raw_json = InferenceProcessor.extract_json_from_output(vlm_output)
320
+
321
+ # Step 5: Validate and fix
322
+ validated_fields, field_confidence, warnings = InferenceProcessor.validate_prediction(raw_json)
323
+
324
+ # Add signature and stamp
325
+ validated_fields["signature"] = signature_info
326
+ validated_fields["stamp"] = stamp_info
327
+
328
+ # Calculate overall confidence
329
+ confidences = [field_confidence]
330
+ if signature_info["present"]:
331
+ confidences.append(signature_conf)
332
+ if stamp_info["present"]:
333
+ confidences.append(stamp_conf)
334
+
335
+ overall_confidence = round(sum(confidences) / len(confidences), 3)
336
+
337
+ # Calculate time and cost
338
+ total_time = time.time() - total_start
339
+ cost_estimate = (COST_PER_GPU_HOUR * total_time) / 3600
340
+
341
+ # Build result
342
+ result = {
343
+ "doc_id": doc_id,
344
+ "fields": validated_fields,
345
+ "confidence": overall_confidence,
346
+ "processing_time_sec": round(total_time, 2),
347
+ "cost_estimate_usd": round(cost_estimate, 6),
348
+ "warnings": warnings if warnings else None
349
+ }
350
+
351
+ return result
model_manager.py ADDED
@@ -0,0 +1,145 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Model Manager - Handles loading and caching of YOLO and VLM models
3
+ """
4
+
5
+ import torch
6
+ from transformers import (
7
+ Qwen2_5_VLForConditionalGeneration,
8
+ AutoProcessor,
9
+ BitsAndBytesConfig
10
+ )
11
+ from ultralytics import YOLO
12
+ import os
13
+ from typing import Tuple
14
+
15
+ from config import (
16
+ YOLO_MODEL_PATH,
17
+ VLM_MODEL_ID,
18
+ QUANTIZATION_CONFIG,
19
+ YOLO_CONFIDENCE_THRESHOLD
20
+ )
21
+
22
+
23
+ class ModelManager:
24
+ """Singleton class to manage model loading and inference"""
25
+
26
+ _instance = None
27
+ _initialized = False
28
+
29
+ def __new__(cls):
30
+ if cls._instance is None:
31
+ cls._instance = super(ModelManager, cls).__new__(cls)
32
+ return cls._instance
33
+
34
+ def __init__(self):
35
+ if not ModelManager._initialized:
36
+ self.yolo_model = None
37
+ self.vlm_model = None
38
+ self.processor = None
39
+ ModelManager._initialized = True
40
+
41
+ def load_models(self):
42
+ """Load both YOLO and VLM models into memory"""
43
+ print("πŸš€ Starting model loading...")
44
+
45
+ # Load YOLO model
46
+ self.yolo_model = self._load_yolo_model()
47
+
48
+ # Load VLM model
49
+ self.vlm_model, self.processor = self._load_vlm_model()
50
+
51
+ print("βœ… All models loaded successfully!")
52
+
53
+ def _load_yolo_model(self) -> YOLO:
54
+ """Load trained YOLO model for signature and stamp detection"""
55
+ if not os.path.exists(YOLO_MODEL_PATH):
56
+ raise FileNotFoundError(
57
+ f"YOLO model not found at {YOLO_MODEL_PATH}. "
58
+ "Please ensure best.pt is in utils/models/"
59
+ )
60
+
61
+ yolo_model = YOLO(str(YOLO_MODEL_PATH))
62
+ print(f"βœ… YOLO model loaded from {YOLO_MODEL_PATH}")
63
+ return yolo_model
64
+
65
+ def _load_vlm_model(self) -> Tuple:
66
+ """
67
+ Load Qwen2.5-VL model with 4-bit quantization
68
+ Downloads from Hugging Face on first run
69
+ """
70
+ print(f"πŸ“₯ Loading VLM model: {VLM_MODEL_ID}")
71
+ print(" (This will download ~4GB on first run)")
72
+
73
+ # Configure 4-bit quantization
74
+ bnb_config = BitsAndBytesConfig(
75
+ load_in_4bit=QUANTIZATION_CONFIG["load_in_4bit"],
76
+ bnb_4bit_quant_type=QUANTIZATION_CONFIG["bnb_4bit_quant_type"],
77
+ bnb_4bit_compute_dtype=getattr(torch, QUANTIZATION_CONFIG["bnb_4bit_compute_dtype"]),
78
+ bnb_4bit_use_double_quant=QUANTIZATION_CONFIG["bnb_4bit_use_double_quant"]
79
+ )
80
+
81
+ # Load processor
82
+ processor = AutoProcessor.from_pretrained(
83
+ VLM_MODEL_ID,
84
+ trust_remote_code=True
85
+ )
86
+
87
+ # Load model with quantization
88
+ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
89
+ VLM_MODEL_ID,
90
+ quantization_config=bnb_config,
91
+ device_map="auto",
92
+ torch_dtype=torch.bfloat16,
93
+ trust_remote_code=True
94
+ )
95
+
96
+ model.eval()
97
+ print(f"βœ… Qwen2.5-VL model loaded successfully")
98
+
99
+ return model, processor
100
+
101
+ def detect_sign_stamp(self, image_path: str):
102
+ """
103
+ Detect signature and stamp in the image using YOLO
104
+
105
+ Returns:
106
+ tuple: (signature_info, stamp_info, signature_conf, stamp_conf)
107
+ """
108
+ if self.yolo_model is None:
109
+ raise RuntimeError("YOLO model not loaded. Call load_models() first.")
110
+
111
+ results = self.yolo_model(image_path, verbose=False)[0]
112
+
113
+ signature_info = {"present": False, "bbox": None}
114
+ stamp_info = {"present": False, "bbox": None}
115
+ signature_conf = 0.0
116
+ stamp_conf = 0.0
117
+
118
+ if results.boxes is not None:
119
+ for box in results.boxes:
120
+ cls_id = int(box.cls[0])
121
+ conf = float(box.conf[0])
122
+
123
+ if conf > YOLO_CONFIDENCE_THRESHOLD:
124
+ bbox = box.xyxy[0].cpu().numpy().tolist()
125
+ bbox = [int(coord) for coord in bbox]
126
+
127
+ # Class 0: signature, Class 1: stamp
128
+ if cls_id == 0 and conf > signature_conf:
129
+ signature_info = {"present": True, "bbox": bbox}
130
+ signature_conf = conf
131
+ elif cls_id == 1 and conf > stamp_conf:
132
+ stamp_info = {"present": True, "bbox": bbox}
133
+ stamp_conf = conf
134
+
135
+ return signature_info, stamp_info, signature_conf, stamp_conf
136
+
137
+ def is_loaded(self) -> bool:
138
+ """Check if models are loaded"""
139
+ return (self.yolo_model is not None and
140
+ self.vlm_model is not None and
141
+ self.processor is not None)
142
+
143
+
144
+ # Global model manager instance
145
+ model_manager = ModelManager()
requirements.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ transformers
3
+ ultralytics
4
+ pillow
5
+ accelerate
6
+ bitsandbytes
7
+ opencv-python
8
+ pyyaml
9
+ qwen-vl-utils[decord]
10
+ fastapi
11
+ uvicorn[standard]
12
+ python-multipart
start.sh ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+
3
+ # Startup script for Invoice Information Extractor API
4
+ # For Hugging Face Spaces or production deployment
5
+
6
+ echo "πŸš€ Invoice Information Extractor - Starting..."
7
+ echo "=============================================="
8
+
9
+ # Check Python version
10
+ python_version=$(python --version 2>&1)
11
+ echo "Python: $python_version"
12
+
13
+ # Check CUDA availability
14
+ if command -v nvidia-smi &> /dev/null; then
15
+ echo "GPU: Available"
16
+ nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
17
+ else
18
+ echo "⚠️ WARNING: No GPU detected. This application requires GPU!"
19
+ fi
20
+
21
+ echo ""
22
+ echo "Starting FastAPI server..."
23
+ echo "=============================================="
24
+
25
+ # Start the application
26
+ python app.py
utils/models/best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:47241e8bafe01e99fc875e1010191d0aa40d797b5f1a5d5110d1cce8b6da8f3d
3
+ size 22508131