Spaces:
Sleeping
Sleeping
Commit ·
35ae7df
1
Parent(s): 7914701
removed print
Browse files
app.py
CHANGED
|
@@ -40,16 +40,12 @@ def load_model(model_choice):
|
|
| 40 |
|
| 41 |
# Only load if different model is requested
|
| 42 |
if current_model != selected_model:
|
| 43 |
-
print(f"Loading model: {selected_model}")
|
| 44 |
try:
|
| 45 |
loaded_generator = pipeline("text2text-generation", model=selected_model)
|
| 46 |
current_model = selected_model
|
| 47 |
-
print(f"Model {selected_model} loaded successfully")
|
| 48 |
except Exception as e:
|
| 49 |
-
print(f"Error loading model {selected_model}: {e}")
|
| 50 |
# Fallback to base model
|
| 51 |
if selected_model != "google/flan-t5-base":
|
| 52 |
-
print("Falling back to base model...")
|
| 53 |
loaded_generator = pipeline("text2text-generation", model="google/flan-t5-base")
|
| 54 |
current_model = "google/flan-t5-base"
|
| 55 |
else:
|
|
@@ -74,11 +70,8 @@ def extract_dl_info(image, model_choice, manual_text, use_manual_text):
|
|
| 74 |
# Determine text source
|
| 75 |
if use_manual_text and manual_text and manual_text.strip():
|
| 76 |
# Use manual text input
|
| 77 |
-
print("Using manual text input...")
|
| 78 |
raw_text = manual_text.strip()
|
| 79 |
cleaned_text = clean_ocr_text(raw_text)
|
| 80 |
-
print(f"Manual text: {raw_text}")
|
| 81 |
-
print(f"Cleaned text: {cleaned_text}")
|
| 82 |
else:
|
| 83 |
# Use OCR on image
|
| 84 |
if image is None:
|
|
@@ -89,23 +82,18 @@ def extract_dl_info(image, model_choice, manual_text, use_manual_text):
|
|
| 89 |
image = image.convert("RGB")
|
| 90 |
|
| 91 |
# Use Tesseract OCR to extract text
|
| 92 |
-
print("Extracting text with Tesseract OCR...")
|
| 93 |
raw_text = pytesseract.image_to_string(image)
|
| 94 |
-
print(f"Raw OCR text: {raw_text}")
|
| 95 |
|
| 96 |
# Clean the OCR text using regex (only for cleaning, not extraction)
|
| 97 |
cleaned_text = clean_ocr_text(raw_text)
|
| 98 |
-
print(f"Cleaned OCR text: {cleaned_text}")
|
| 99 |
|
| 100 |
# Use AI model to process cleaned OCR text and extract structured information
|
| 101 |
-
print("Processing cleaned OCR text with AI model...")
|
| 102 |
|
| 103 |
# Load model only when needed
|
| 104 |
try:
|
| 105 |
generator = load_model(model_choice)
|
| 106 |
except Exception as e:
|
| 107 |
error_msg = f"Failed to load AI model: {str(e)}"
|
| 108 |
-
print(error_msg)
|
| 109 |
return raw_text, cleaned_text, json.dumps({"error": error_msg}, indent=2)
|
| 110 |
|
| 111 |
# Create specific prompts for the AI model to extract individual fields
|
|
@@ -156,7 +144,6 @@ def extract_dl_info(image, model_choice, manual_text, use_manual_text):
|
|
| 156 |
}
|
| 157 |
|
| 158 |
except Exception as e:
|
| 159 |
-
print(f"AI extraction failed: {e}")
|
| 160 |
extracted = {
|
| 161 |
"name": "AI extraction failed",
|
| 162 |
"address": "AI extraction failed",
|
|
|
|
| 40 |
|
| 41 |
# Only load if different model is requested
|
| 42 |
if current_model != selected_model:
|
|
|
|
| 43 |
try:
|
| 44 |
loaded_generator = pipeline("text2text-generation", model=selected_model)
|
| 45 |
current_model = selected_model
|
|
|
|
| 46 |
except Exception as e:
|
|
|
|
| 47 |
# Fallback to base model
|
| 48 |
if selected_model != "google/flan-t5-base":
|
|
|
|
| 49 |
loaded_generator = pipeline("text2text-generation", model="google/flan-t5-base")
|
| 50 |
current_model = "google/flan-t5-base"
|
| 51 |
else:
|
|
|
|
| 70 |
# Determine text source
|
| 71 |
if use_manual_text and manual_text and manual_text.strip():
|
| 72 |
# Use manual text input
|
|
|
|
| 73 |
raw_text = manual_text.strip()
|
| 74 |
cleaned_text = clean_ocr_text(raw_text)
|
|
|
|
|
|
|
| 75 |
else:
|
| 76 |
# Use OCR on image
|
| 77 |
if image is None:
|
|
|
|
| 82 |
image = image.convert("RGB")
|
| 83 |
|
| 84 |
# Use Tesseract OCR to extract text
|
|
|
|
| 85 |
raw_text = pytesseract.image_to_string(image)
|
|
|
|
| 86 |
|
| 87 |
# Clean the OCR text using regex (only for cleaning, not extraction)
|
| 88 |
cleaned_text = clean_ocr_text(raw_text)
|
|
|
|
| 89 |
|
| 90 |
# Use AI model to process cleaned OCR text and extract structured information
|
|
|
|
| 91 |
|
| 92 |
# Load model only when needed
|
| 93 |
try:
|
| 94 |
generator = load_model(model_choice)
|
| 95 |
except Exception as e:
|
| 96 |
error_msg = f"Failed to load AI model: {str(e)}"
|
|
|
|
| 97 |
return raw_text, cleaned_text, json.dumps({"error": error_msg}, indent=2)
|
| 98 |
|
| 99 |
# Create specific prompts for the AI model to extract individual fields
|
|
|
|
| 144 |
}
|
| 145 |
|
| 146 |
except Exception as e:
|
|
|
|
| 147 |
extracted = {
|
| 148 |
"name": "AI extraction failed",
|
| 149 |
"address": "AI extraction failed",
|