solarplant / app.py
npc0's picture
Create app.py
79f9dfb
raw
history blame
1.17 kB
import os
os.system('wget https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/FastSAM.pt')
import yolov5
# load model
model = yolov5.load('keremberke/yolov5m-license-plate')
# set model parameters
model.conf = 0.5 # NMS confidence threshold
model.iou = 0.25 # NMS IoU threshold
model.agnostic = False # NMS class-agnostic
model.multi_label = False # NMS multiple labels per box
model.max_det = 1000 # maximum number of detections per image
# set image
def license_plate_detect(img):
# perform inference
results = model(img, size=640)
# inference with test time augmentation
results = model(img, augment=True)
# parse results
if len(results.pred):
predictions = results.pred[0]
boxes = predictions[:, :4] # x1, y1, x2, y2
scores = predictions[:, 4]
categories = predictions[:, 5]
return boxes
from PIL import Image
# image = Image.open(img)
import pytesseract
def read_license_number(img):
boxes = license_plate_detect(img)
if boxes:
return [pytesseract.image_to_string(
image.crop(bbox.tolist()))
for bbox in boxes]