Rammohan0504 commited on
Commit
e32d5ab
·
verified ·
1 Parent(s): 5d52c44

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -7
app.py CHANGED
@@ -11,7 +11,7 @@ from transformers import TrOCRProcessor, VisionEncoderDecoderModel
11
  # Load models
12
  device = "cuda" if torch.cuda.is_available() else "cpu"
13
 
14
- # YOLOv5 for digit detection (Pre-trained model)
15
  yolo_model = YOLO("yolov5s.pt")
16
 
17
  # OCR Models
@@ -33,18 +33,18 @@ def convert_to_rgb(image):
33
  image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
34
  return image
35
 
36
- # Detect Digits Using YOLOv5
37
- def detect_digits(image):
38
  results = yolo_model(image)
39
- detected_numbers = []
40
 
41
  for result in results:
42
  if hasattr(result, "boxes"): # Ensure correct format
43
  for box in result.boxes:
44
  if box.conf > 0.5: # Confidence threshold
45
- detected_numbers.append(box.xyxy.tolist())
46
 
47
- return detected_numbers
48
 
49
  # Extract Text Using EasyOCR
50
  def extract_text_easyocr(image):
@@ -68,7 +68,7 @@ def extract_weight(text):
68
  # Full Processing Pipeline
69
  def process_image(image):
70
  enhanced = enhance_image(image)
71
- detected_digits = detect_digits(image)
72
  text_easyocr = extract_text_easyocr(enhanced)
73
  text_trocr = extract_text_trocr(enhanced)
74
 
 
11
  # Load models
12
  device = "cuda" if torch.cuda.is_available() else "cpu"
13
 
14
+ # YOLOv5 for digital meter detection (Pre-trained model)
15
  yolo_model = YOLO("yolov5s.pt")
16
 
17
  # OCR Models
 
33
  image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
34
  return image
35
 
36
+ # Detect Digital Meter Using YOLOv5
37
+ def detect_meter(image):
38
  results = yolo_model(image)
39
+ detected_meters = []
40
 
41
  for result in results:
42
  if hasattr(result, "boxes"): # Ensure correct format
43
  for box in result.boxes:
44
  if box.conf > 0.5: # Confidence threshold
45
+ detected_meters.append(box.xyxy.tolist())
46
 
47
+ return detected_meters
48
 
49
  # Extract Text Using EasyOCR
50
  def extract_text_easyocr(image):
 
68
  # Full Processing Pipeline
69
  def process_image(image):
70
  enhanced = enhance_image(image)
71
+ detected_meters = detect_meter(image)
72
  text_easyocr = extract_text_easyocr(enhanced)
73
  text_trocr = extract_text_trocr(enhanced)
74