Spaces:
Sleeping
Sleeping
| import pytesseract | |
| import numpy as np | |
| import cv2 | |
| import re | |
| import logging | |
| from datetime import datetime | |
| import os | |
| from PIL import Image | |
| # Set up logging | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| # Directory for debug images | |
| DEBUG_DIR = "debug_images" | |
| os.makedirs(DEBUG_DIR, exist_ok=True) | |
| def save_debug_image(img, filename_suffix, prefix=""): | |
| """Save image to debug directory with timestamp.""" | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") | |
| filename = os.path.join(DEBUG_DIR, f"{prefix}{timestamp}_{filename_suffix}.png") | |
| if isinstance(img, Image.Image): | |
| img.save(filename) | |
| elif len(img.shape) == 3: | |
| cv2.imwrite(filename, cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) | |
| else: | |
| cv2.imwrite(filename, img) | |
| logging.info(f"Saved debug image: {filename}") | |
| def estimate_brightness(img): | |
| """Estimate image brightness.""" | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| return np.mean(gray) | |
| def preprocess_image(img): | |
| """Preprocess image with simplified, robust contrast enhancement.""" | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| brightness = estimate_brightness(img) | |
| # Apply mild CLAHE for contrast | |
| clahe_clip = 8.0 if brightness < 90 else 4.0 | |
| clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=(8, 8)) | |
| enhanced = clahe.apply(gray) | |
| save_debug_image(enhanced, "01_preprocess_clahe") | |
| # Light blur to reduce noise | |
| blurred = cv2.GaussianBlur(enhanced, (5, 5), 0) | |
| save_debug_image(blurred, "02_preprocess_blur") | |
| # Dynamic thresholding with larger block size for small displays | |
| block_size = max(7, min(31, int(img.shape[0] / 20) * 2 + 1)) | |
| thresh = cv2.adaptiveThreshold( | |
| blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, | |
| cv2.THRESH_BINARY_INV, block_size, 3 | |
| ) | |
| # Minimal morphological operations | |
| kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) | |
| thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1) | |
| save_debug_image(thresh, "03_preprocess_morph") | |
| return thresh, enhanced | |
| def correct_rotation(img): | |
| """Correct image rotation using edge detection.""" | |
| try: | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| edges = cv2.Canny(gray, 30, 100, apertureSize=3) | |
| lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=25, minLineLength=15, maxLineGap=10) | |
| if lines is not None: | |
| angles = [np.arctan2(line[0][3] - line[0][1], line[0][2] - line[0][0]) * 180 / np.pi for line in lines] | |
| angle = np.median(angles) | |
| if abs(angle) > 0.3: | |
| h, w = img.shape[:2] | |
| center = (w // 2, h // 2) | |
| M = cv2.getRotationMatrix2D(center, angle, 1.0) | |
| img = cv2.warpAffine(img, M, (w, h)) | |
| save_debug_image(img, "00_rotated_image") | |
| logging.info(f"Applied rotation: {angle:.2f} degrees") | |
| return img | |
| except Exception as e: | |
| logging.error(f"Rotation correction failed: {str(e)}") | |
| return img | |
| def detect_roi(img): | |
| """Detect region of interest with broader contour analysis.""" | |
| try: | |
| save_debug_image(img, "04_original") | |
| thresh, enhanced = preprocess_image(img) | |
| brightness_map = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| block_sizes = [max(7, min(31, int(img.shape[0] / s) * 2 + 1)) for s in [5, 10, 20]] | |
| valid_contours = [] | |
| img_area = img.shape[0] * img.shape[1] | |
| for block_size in block_sizes: | |
| temp_thresh = cv2.adaptiveThreshold( | |
| enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, | |
| cv2.THRESH_BINARY_INV, block_size, 3 | |
| ) | |
| kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) | |
| temp_thresh = cv2.morphologyEx(temp_thresh, cv2.MORPH_CLOSE, kernel, iterations=2) | |
| save_debug_image(temp_thresh, f"05_roi_threshold_block{block_size}") | |
| contours, _ = cv2.findContours(temp_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| for c in contours: | |
| area = cv2.contourArea(c) | |
| x, y, w, h = cv2.boundingRect(c) | |
| roi_brightness = np.mean(brightness_map[y:y+h, x:x+w]) | |
| aspect_ratio = w / h | |
| if (50 < area < (img_area * 0.95) and | |
| 0.05 <= aspect_ratio <= 20.0 and w > 20 and h > 8 and roi_brightness > 15): | |
| valid_contours.append((c, area * roi_brightness)) | |
| logging.debug(f"Contour (block {block_size}): Area={area}, Aspect={aspect_ratio:.2f}, Brightness={roi_brightness:.2f}") | |
| if valid_contours: | |
| contour, _ = max(valid_contours, key=lambda x: x[1]) | |
| x, y, w, h = cv2.boundingRect(contour) | |
| padding = max(5, min(20, int(min(w, h) * 0.4))) | |
| x, y = max(0, x - padding), max(0, y - padding) | |
| w, h = min(w + 2 * padding, img.shape[1] - x), min(h + 2 * padding, img.shape[0] - y) | |
| roi_img = img[y:y+h, x:x+w] | |
| save_debug_image(roi_img, "06_detected_roi") | |
| logging.info(f"Detected ROI: ({x}, {y}, {w}, {h})") | |
| return roi_img, (x, y, w, h) | |
| logging.info("No ROI found, using full image.") | |
| save_debug_image(img, "06_no_roi_fallback") | |
| return img, None | |
| except Exception as e: | |
| logging.error(f"ROI detection failed: {str(e)}") | |
| save_debug_image(img, "06_roi_error_fallback") | |
| return img, None | |
| def detect_digit_template(digit_img, brightness): | |
| """Digit recognition with expanded template matching.""" | |
| try: | |
| h, w = digit_img.shape | |
| if h < 5 or w < 2: | |
| logging.debug("Digit image too small for template matching.") | |
| return None | |
| # Expanded digit templates for seven-segment display variations | |
| digit_templates = { | |
| '0': [ | |
| np.array([[1, 1, 1, 1, 1], | |
| [1, 0, 0, 0, 1], | |
| [1, 0, 0, 0, 1], | |
| [1, 0, 0, 0, 1], | |
| [1, 1, 1, 1, 1]], dtype=np.float32), | |
| np.array([[1, 1, 1, 1], | |
| [1, 0, 0, 1], | |
| [1, 0, 0, 1], | |
| [1, 0, 0, 1], | |
| [1, 1, 1, 1]], dtype=np.float32) | |
| ], | |
| '1': [ | |
| np.array([[0, 0, 1, 0, 0], | |
| [0, 0, 1, 0, 0], | |
| [0, 0, 1, 0, 0], | |
| [0, 0, 1, 0, 0], | |
| [0, 0, 1, 0, 0]], dtype=np.float32), | |
| np.array([[0, 1, 0], | |
| [0, 1, 0], | |
| [0, 1, 0], | |
| [0, 1, 0], | |
| [0, 1, 0]], dtype=np.float32) | |
| ], | |
| '2': [ | |
| np.array([[1, 1, 1, 1, 1], | |
| [0, 0, 0, 1, 1], | |
| [1, 1, 1, 1, 1], | |
| [1, 1, 0, 0, 0], | |
| [1, 1, 1, 1, 1]], dtype=np.float32), | |
| np.array([[1, 1, 1, 1], | |
| [0, 0, 1, 1], | |
| [1, 1, 1, 1], | |
| [1, 1, 0, 0], | |
| [1, 1, 1, 1]], dtype=np.float32) | |
| ], | |
| '3': [ | |
| np.array([[1, 1, 1, 1, 1], | |
| [0, 0, 0, 1, 1], | |
| [1, 1, 1, 1, 1], | |
| [0, 0, 0, 1, 1], | |
| [1, 1, 1, 1, 1]], dtype=np.float32), | |
| np.array([[1, 1, 1, 1], | |
| [0, 0, 1, 1], | |
| [1, 1, 1, 1], | |
| [0, 0, 1, 1], | |
| [1, 1, 1, 1]], dtype=np.float32) | |
| ], | |
| '4': [ | |
| np.array([[1, 1, 0, 0, 1], | |
| [1, 1, 0, 0, 1], | |
| [1, 1, 1, 1, 1], | |
| [0, 0, 0, 0, 1], | |
| [0, 0, 0, 0, 1]], dtype=np.float32), | |
| np.array([[1, 0, 0, 1], | |
| [1, 0, 0, 1], | |
| [1, 1, 1, 1], | |
| [0, 0, 0, 1], | |
| [0, 0, 0, 1]], dtype=np.float32) | |
| ], | |
| '5': [ | |
| np.array([[1, 1, 1, 1, 1], | |
| [1, 1, 0, 0, 0], | |
| [1, 1, 1, 1, 1], | |
| [0, 0, 0, 1, 1], | |
| [1, 1, 1, 1, 1]], dtype=np.float32), | |
| np.array([[1, 1, 1, 1], | |
| [1, 1, 0, 0], | |
| [1, 1, 1, 1], | |
| [0, 0, 1, 1], | |
| [1, 1, 1, 1]], dtype=np.float32) | |
| ], | |
| '6': [ | |
| np.array([[1, 1, 1, 1, 1], | |
| [1, 1, 0, 0, 0], | |
| [1, 1, 1, 1, 1], | |
| [1, 0, 0, 1, 1], | |
| [1, 1, 1, 1, 1]], dtype=np.float32), | |
| np.array([[1, 1, 1, 1], | |
| [1, 1, 0, 0], | |
| [1, 1, 1, 1], | |
| [1, 0, 1, 1], | |
| [1, 1, 1, 1]], dtype=np.float32) | |
| ], | |
| '7': [ | |
| np.array([[1, 1, 1, 1, 1], | |
| [0, 0, 0, 0, 1], | |
| [0, 0, 0, 0, 1], | |
| [0, 0, 0, 0, 1], | |
| [0, 0, 0, 0, 1]], dtype=np.float32), | |
| np.array([[1, 1, 1, 1], | |
| [0, 0, 0, 1], | |
| [0, 0, 0, 1], | |
| [0, 0, 0, 1], | |
| [0, 0, 0, 1]], dtype=np.float32) | |
| ], | |
| '8': [ | |
| np.array([[1, 1, 1, 1, 1], | |
| [1, 0, 0, 0, 1], | |
| [1, 1, 1, 1, 1], | |
| [1, 0, 0, 0, 1], | |
| [1, 1, 1, 1, 1]], dtype=np.float32), | |
| np.array([[1, 1, 1, 1], | |
| [1, 0, 0, 1], | |
| [1, 1, 1, 1], | |
| [1, 0, 0, 1], | |
| [1, 1, 1, 1]], dtype=np.float32) | |
| ], | |
| '9': [ | |
| np.array([[1, 1, 1, 1, 1], | |
| [1, 0, 0, 0, 1], | |
| [1, 1, 1, 1, 1], | |
| [0, 0, 0, 1, 1], | |
| [1, 1, 1, 1, 1]], dtype=np.float32), | |
| np.array([[1, 1, 1, 1], | |
| [1, 0, 0, 1], | |
| [1, 1, 1, 1], | |
| [0, 0, 1, 1], | |
| [1, 1, 1, 1]], dtype=np.float32) | |
| ], | |
| '.': [ | |
| np.array([[0, 0, 0], | |
| [0, 1, 0], | |
| [0, 0, 0]], dtype=np.float32), | |
| np.array([[0, 0], | |
| [1, 0], | |
| [0, 0]], dtype=np.float32) | |
| ] | |
| } | |
| # Try multiple sizes for digit image | |
| sizes = [(5, 5), (4, 4), (3, 3)] if h > w else [(3, 3), (2, 2)] | |
| best_match, best_score = None, -1 | |
| for size in sizes: | |
| digit_img_resized = cv2.resize(digit_img, size, interpolation=cv2.INTER_AREA) | |
| digit_img_resized = (digit_img_resized > 100).astype(np.float32) # Binarize | |
| for digit, templates in digit_templates.items(): | |
| for template in templates: | |
| if digit == '.' and size[0] > 3: | |
| continue | |
| if digit != '.' and size[0] <= 3: | |
| continue | |
| if template.shape[0] != size[0] or template.shape[1] != size[1]: | |
| continue | |
| result = cv2.matchTemplate(digit_img_resized, template, cv2.TM_CCOEFF_NORMED) | |
| _, max_val, _, _ = cv2.minMaxLoc(result) | |
| if max_val > 0.55 and max_val > best_score: # Further lowered threshold | |
| best_score = max_val | |
| best_match = digit | |
| logging.debug(f"Template match: {best_match}, Score: {best_score:.2f}") | |
| return best_match if best_score > 0.55 else None | |
| except Exception as e: | |
| logging.error(f"Template digit detection failed: {str(e)}") | |
| return None | |
| def perform_ocr(img, roi_bbox): | |
| """Perform OCR with Tesseract and robust template fallback.""" | |
| try: | |
| thresh, enhanced = preprocess_image(img) | |
| brightness = estimate_brightness(img) | |
| pil_img = Image.fromarray(enhanced) | |
| save_debug_image(pil_img, "07_ocr_input") | |
| # Try multiple Tesseract configurations | |
| configs = [ | |
| r'--oem 3 --psm 7 -c tessedit_char_whitelist=0123456789.', # Single line | |
| r'--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789.' # Block of text | |
| ] | |
| for config in configs: | |
| text = pytesseract.image_to_string(pil_img, config=config) | |
| logging.info(f"Tesseract raw output (config {config}): {text}") | |
| text = re.sub(r"[^\d\.]", "", text) | |
| if text.count('.') > 1: | |
| text = text.replace('.', '', text.count('.') - 1) | |
| text = text.strip('.') | |
| if text and re.fullmatch(r"^\d*\.?\d*$", text): | |
| text = text.lstrip('0') or '0' | |
| confidence = 95.0 if len(text.replace('.', '')) >= 3 else 90.0 | |
| logging.info(f"Validated Tesseract text: {text}, Confidence: {confidence:.2f}%") | |
| return text, confidence | |
| # Fallback to template-based detection | |
| logging.info("Tesseract failed, using template-based detection.") | |
| contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| digits_info = [] | |
| for c in contours: | |
| x, y, w, h = cv2.boundingRect(c) | |
| if w > 4 and h > 5 and 0.03 <= w/h <= 4.0: | |
| digits_info.append((x, x+w, y, y+h)) | |
| if digits_info: | |
| digits_info.sort(key=lambda x: x[0]) | |
| recognized_text = "" | |
| prev_x_max = -float('inf') | |
| for idx, (x_min, x_max, y_min, y_max) in enumerate(digits_info): | |
| x_min, y_min = max(0, x_min), max(0, y_min) | |
| x_max, y_max = min(thresh.shape[1], x_max), min(thresh.shape[0], y_max) | |
| if x_max <= x_min or y_max <= y_min: | |
| continue | |
| digit_crop = thresh[y_min:y_max, x_min:x_max] | |
| save_debug_image(digit_crop, f"08_digit_crop_{idx}") | |
| digit = detect_digit_template(digit_crop, brightness) | |
| if digit: | |
| recognized_text += digit | |
| elif x_min - prev_x_max < 10 and prev_x_max != -float('inf'): | |
| recognized_text += '.' | |
| prev_x_max = x_max | |
| text = re.sub(r"[^\d\.]", "", recognized_text) | |
| if text.count('.') > 1: | |
| text = text.replace('.', '', text.count('.') - 1) | |
| text = text.strip('.') | |
| if text and re.fullmatch(r"^\d*\.?\d*$", text): | |
| text = text.lstrip('0') or '0' | |
| confidence = 90.0 if len(text.replace('.', '')) >= 3 else 85.0 | |
| logging.info(f"Validated template text: {text}, Confidence: {confidence:.2f}%") | |
| return text, confidence | |
| logging.info("No valid digits detected.") | |
| return None, 0.0 | |
| except Exception as e: | |
| logging.error(f"OCR failed: {str(e)}") | |
| return None, 0.0 | |
| def extract_weight_from_image(pil_img): | |
| """Extract weight from any digital scale image.""" | |
| try: | |
| img = np.array(pil_img) | |
| img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) | |
| save_debug_image(img, "00_input_image") | |
| img = correct_rotation(img) | |
| brightness = estimate_brightness(img) | |
| conf_threshold = 0.65 if brightness > 70 else 0.45 | |
| # Try ROI-based detection | |
| roi_img, roi_bbox = detect_roi(img) | |
| if roi_bbox: | |
| conf_threshold *= 1.15 if (roi_bbox[2] * roi_bbox[3]) > (img.shape[0] * img.shape[1] * 0.05) else 1.0 | |
| result, confidence = perform_ocr(roi_img, roi_bbox) | |
| if result and confidence >= conf_threshold * 100: | |
| try: | |
| weight = float(result) | |
| if 0.001 <= weight <= 5000: | |
| logging.info(f"Detected weight: {result} kg, Confidence: {confidence:.2f}%") | |
| return result, confidence | |
| logging.warning(f"Weight {result} out of range.") | |
| except ValueError: | |
| logging.warning(f"Invalid weight format: {result}") | |
| # Full image fallback | |
| logging.info("Primary OCR failed, using full image fallback.") | |
| result, confidence = perform_ocr(img, None) | |
| if result and confidence >= conf_threshold * 0.85 * 100: | |
| try: | |
| weight = float(result) | |
| if 0.001 <= weight <= 5000: | |
| logging.info(f"Full image weight: {result} kg, Confidence: {confidence:.2f}%") | |
| return result, confidence | |
| logging.warning(f"Full image weight {result} out of range.") | |
| except ValueError: | |
| logging.warning(f"Invalid full image weight format: {result}") | |
| logging.info("No valid weight detected.") | |
| return "Not detected", 0.0 | |
| except Exception as e: | |
| logging.error(f"Weight extraction failed: {str(e)}") | |
| return "Not detected", 0.0 |