Upload app.py with huggingface_hub

app.py

@@ -7,15 +7,22 @@ from ultralytics import YOLO
 max_len = 10
 
 def ocr(plate_image, ocr_model):
-    # This function should return the OCR result of a license plate image
-    # For demo purposes, let's return a dummy value
-    return "ABC1234"
-
+  try:
+    if 'yolo' in ocr_model:
+      model = YOLO(f'/content/drive/MyDrive/{ocr_model}.pt', task='detect')
+      preds = model.predict(plate_image)
+      return ''.join([model.names[int(pred.boxes.cls.cpu().numpy()[0])] for pred in preds[0]])
+    else:
+      return ''
+  except Exception : 
+    
+    return 'error'
 
 def process_image(image, detection_model, ocr_model, yolo_thresh, perpective_width, perpective_height):
-    model = YOLO(f'./weights/{detection_model}.pt')
     task = 'obb' if 'obb' in detection_model else 'detect'
-    predict  = model(image, task=task, conf=yolo_thresh)
+    model = YOLO(f'/content/drive/MyDrive/{detection_model}.pt', task=task)
+
+    predict  = model(image, conf=yolo_thresh)
     
     if 'obb' in detection_model:
       obb_crops, transformed = perspective_transform(predict, dst_width=perpective_width, dst_height=perpective_height)  
@@ -37,7 +44,7 @@ def create_interface():
         with gr.Row():
             with gr.Column(scale=1):
               detection_model = gr.Dropdown(label="Detection Model", choices=["yolov8-m", "yolov8-obb-m","yolov8-s", "yolov8-obb-s"], value="yolov8-m")
-              ocr_model = gr.Dropdown(label="OCR Model", choices=["yolov8-32class", "trocr"], value="trocr")
+              ocr_model = gr.Dropdown(label="OCR Model", choices=["trocr-v1", "trocr-v2"], value="trocr-v1")
               yolo_thresh = gr.Slider(minimum=0.0, maximum=1.0, label='yolo_threshold')
               with gr.Row():
                 perpective_width = gr.Slider(minimum=100, maximum=500, label='perpective_width')