Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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
|
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
|
37 |
-
def
|
38 |
results = yolo_model(image)
|
39 |
-
|
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 |
-
|
46 |
|
47 |
-
return
|
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 |
-
|
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 |
|