Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -20,7 +20,7 @@ ocr_reader = easyocr.Reader(["en"]) # EasyOCR
|
|
20 |
trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
|
21 |
trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1").to(device)
|
22 |
|
23 |
-
# Image Preprocessing (
|
24 |
def enhance_image(image):
|
25 |
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
26 |
|
@@ -40,7 +40,7 @@ def convert_to_rgb(image):
|
|
40 |
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
|
41 |
return image
|
42 |
|
43 |
-
# Detect Digital Meter Using YOLOv5
|
44 |
def detect_meter(image):
|
45 |
results = yolo_model(image)
|
46 |
detected_meters = []
|
@@ -48,7 +48,7 @@ def detect_meter(image):
|
|
48 |
for result in results:
|
49 |
if hasattr(result, "boxes"): # Ensure correct format
|
50 |
for box in result.boxes:
|
51 |
-
if box.conf > 0.
|
52 |
detected_meters.append(box.xyxy.tolist())
|
53 |
|
54 |
return detected_meters
|
@@ -80,10 +80,12 @@ def filter_weight_values(weights):
|
|
80 |
return str(round(filtered_weight, 2))
|
81 |
return weights[0] if weights else "Weight not detected"
|
82 |
|
83 |
-
# Full Processing Pipeline
|
84 |
def process_image(image):
|
85 |
enhanced = enhance_image(image)
|
86 |
detected_meters = detect_meter(image)
|
|
|
|
|
87 |
text_easyocr = extract_text_easyocr(enhanced)
|
88 |
text_trocr = extract_text_trocr(enhanced)
|
89 |
|
@@ -94,7 +96,11 @@ def process_image(image):
|
|
94 |
final_weights = [weight_easyocr, weight_trocr]
|
95 |
final_weight = filter_weight_values([w for w in final_weights if w != "Weight not detected"])
|
96 |
|
97 |
-
|
|
|
|
|
|
|
|
|
98 |
|
99 |
# Gradio Interface
|
100 |
iface = gr.Interface(fn=process_image, inputs="image", outputs="text")
|
|
|
20 |
trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
|
21 |
trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1").to(device)
|
22 |
|
23 |
+
# Image Preprocessing (Adaptive Threshold & Sharpening)
|
24 |
def enhance_image(image):
|
25 |
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
26 |
|
|
|
40 |
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
|
41 |
return image
|
42 |
|
43 |
+
# Detect Digital Meter Using YOLOv5 (Improved Confidence Threshold)
|
44 |
def detect_meter(image):
|
45 |
results = yolo_model(image)
|
46 |
detected_meters = []
|
|
|
48 |
for result in results:
|
49 |
if hasattr(result, "boxes"): # Ensure correct format
|
50 |
for box in result.boxes:
|
51 |
+
if box.conf > 0.25: # Lower confidence threshold for better detection
|
52 |
detected_meters.append(box.xyxy.tolist())
|
53 |
|
54 |
return detected_meters
|
|
|
80 |
return str(round(filtered_weight, 2))
|
81 |
return weights[0] if weights else "Weight not detected"
|
82 |
|
83 |
+
# Full Processing Pipeline (With OCR Fallback)
|
84 |
def process_image(image):
|
85 |
enhanced = enhance_image(image)
|
86 |
detected_meters = detect_meter(image)
|
87 |
+
|
88 |
+
# OCR Extraction
|
89 |
text_easyocr = extract_text_easyocr(enhanced)
|
90 |
text_trocr = extract_text_trocr(enhanced)
|
91 |
|
|
|
96 |
final_weights = [weight_easyocr, weight_trocr]
|
97 |
final_weight = filter_weight_values([w for w in final_weights if w != "Weight not detected"])
|
98 |
|
99 |
+
# Handle failed detection cases
|
100 |
+
if final_weight == "Weight not detected":
|
101 |
+
return "Try adjusting image clarity or detection thresholds"
|
102 |
+
|
103 |
+
return final_weight
|
104 |
|
105 |
# Gradio Interface
|
106 |
iface = gr.Interface(fn=process_image, inputs="image", outputs="text")
|