Rammohan0504 commited on
Commit
166e55c
·
verified ·
1 Parent(s): 71a7d42

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -5
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 (Sharpen & Adaptive Threshold)
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.3: # Lower confidence threshold for better detection
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
- return final_weight or "Weight not detected"
 
 
 
 
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")