Update app.py

app.py

@@ -12,7 +12,7 @@ from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # YOLOv5 for digit detection (Pre-trained model)
-yolo_model = YOLO("yolov5s.pt")  
+yolo_model = YOLO("yolov5s.pt")
 
 # OCR Models
 ocr_reader = easyocr.Reader(["en"])  # EasyOCR
@@ -30,7 +30,14 @@ def enhance_image(image):
 # Detect Digits Using YOLOv5
 def detect_digits(image):
     results = yolo_model(image)
-    detected_numbers = [det.xyxy.tolist()[0] for det in results.pred[0] if det.conf > 0.5]
+    detected_numbers = []
+
+    for result in results:
+        if hasattr(result, "boxes"):  # Ensure correct format
+            for box in result.boxes:
+                if box.conf > 0.5:  # Confidence threshold
+                    detected_numbers.append(box.xyxy.tolist())
+
     return detected_numbers
 
 # Extract Text Using EasyOCR
@@ -57,11 +64,11 @@ def process_image(image):
     detected_digits = detect_digits(image)
     text_easyocr = extract_text_easyocr(enhanced)
     text_trocr = extract_text_trocr(enhanced)
-    
+
     # Prioritize numeric values from OCR
     weight_easyocr = extract_weight(text_easyocr)
     weight_trocr = extract_weight(text_trocr)
-    
+
     final_weight = weight_easyocr if weight_easyocr != "Weight not detected" else weight_trocr
     return final_weight or "Weight not detected"