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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -1
app.py CHANGED
@@ -27,6 +27,12 @@ def enhance_image(image):
27
  _, thresholded = cv2.threshold(sharpened, 150, 255, cv2.THRESH_BINARY)
28
  return thresholded
29
 
 
 
 
 
 
 
30
  # Detect Digits Using YOLOv5
31
  def detect_digits(image):
32
  results = yolo_model(image)
@@ -45,8 +51,9 @@ def extract_text_easyocr(image):
45
  text = " ".join(ocr_reader.readtext(image, detail=0))
46
  return text
47
 
48
- # Extract Text Using TrOCR
49
  def extract_text_trocr(image):
 
50
  image = Image.fromarray(image)
51
  pixel_values = trocr_processor(images=image, return_tensors="pt").pixel_values.to(device)
52
  generated_ids = trocr_model.generate(pixel_values)
 
27
  _, thresholded = cv2.threshold(sharpened, 150, 255, cv2.THRESH_BINARY)
28
  return thresholded
29
 
30
+ # Convert Grayscale to RGB (Fix for TrOCR)
31
+ def convert_to_rgb(image):
32
+ if len(image.shape) == 2: # Grayscale 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)
 
51
  text = " ".join(ocr_reader.readtext(image, detail=0))
52
  return text
53
 
54
+ # Extract Text Using TrOCR (Fixed)
55
  def extract_text_trocr(image):
56
+ image = convert_to_rgb(image) # Convert grayscale to RGB
57
  image = Image.fromarray(image)
58
  pixel_values = trocr_processor(images=image, return_tensors="pt").pixel_values.to(device)
59
  generated_ids = trocr_model.generate(pixel_values)