Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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)
|