Rammohan0504 commited on
Commit
2b25983
·
verified ·
1 Parent(s): 0a41ce0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -19
app.py CHANGED
@@ -1,17 +1,15 @@
1
- import cv2
2
- import numpy as np
3
  import torch
4
- import pytesseract
5
- import easyocr
6
  import gradio as gr
7
- from yolov5 import YOLOv5
8
  import re
 
 
9
 
10
- # Load YOLOv5 model (pre-trained)
11
- model = YOLOv5('yolov5s.pt', device='cpu')
12
-
13
- # Load EasyOCR
14
- ocr_reader = easyocr.Reader(['en'])
15
 
16
  # Image Preprocessing (Sharpen & Deblur)
17
  def enhance_image(image):
@@ -20,15 +18,12 @@ def enhance_image(image):
20
  sharpened = cv2.filter2D(image, -1, kernel)
21
  return sharpened
22
 
23
- # Detect Digits Using YOLOv5
24
- def detect_digits(image):
25
- results = model(image)
26
- digits = [det.xyxy.tolist()[0] for det in results.pred[0] if det.conf > 0.5]
27
- return digits
28
-
29
- # Extract Text Using OCR
30
  def extract_text(image):
31
- text = pytesseract.image_to_string(image, config='--psm 6')
 
 
 
32
  return text
33
 
34
  # Extract Weight Using Regex
@@ -39,7 +34,6 @@ def extract_weight(text):
39
  # Full Processing Pipeline
40
  def process_image(image):
41
  enhanced = enhance_image(image)
42
- digits = detect_digits(image)
43
  text = extract_text(enhanced)
44
  weight = extract_weight(text)
45
  return weight or "No weight detected"
 
 
 
1
  import torch
2
+ from transformers import BlipProcessor, BlipForConditionalGeneration
3
+ from PIL import Image
4
  import gradio as gr
 
5
  import re
6
+ import cv2
7
+ import numpy as np
8
 
9
+ # Load BLIP model
10
+ device = "cuda" if torch.cuda.is_available() else "cpu"
11
+ processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
12
+ model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
 
13
 
14
  # Image Preprocessing (Sharpen & Deblur)
15
  def enhance_image(image):
 
18
  sharpened = cv2.filter2D(image, -1, kernel)
19
  return sharpened
20
 
21
+ # Extract Text Using BLIP
 
 
 
 
 
 
22
  def extract_text(image):
23
+ image = Image.fromarray(image)
24
+ inputs = processor(images=image, return_tensors="pt").to(device)
25
+ generated_ids = model.generate(**inputs)
26
+ text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
27
  return text
28
 
29
  # Extract Weight Using Regex
 
34
  # Full Processing Pipeline
35
  def process_image(image):
36
  enhanced = enhance_image(image)
 
37
  text = extract_text(enhanced)
38
  weight = extract_weight(text)
39
  return weight or "No weight detected"