Update app.py

app.py

@@ -4,11 +4,18 @@ import time
 import torch
 import argparse
 import gradio as gr
+import io
 from PIL import Image
 from numpy import random
 from pathlib import Path
 import torch.backends.cudnn as cudnn
 from models.experimental import attempt_load
+import keras_ocr
+import matplotlib.pyplot as plt
+from numpy import asarray
+import pytesseract
+from datetime import date
+
 
 from utils.datasets import LoadStreams, LoadImages
 from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
@@ -18,14 +25,13 @@ from utils.torch_utils import select_device, load_classifier, time_synchronized,
 os.system("wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt")
 os.system("wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-e6.pt")
 
-def detect_Custom(img,model):
-    if model =='Yolo_v7_Custom_trained_By_Owais':
-        model='best' # Naming Convention for yolov7 See output file of https://www.kaggle.com/code/owaiskhan9654/training-yolov7-on-kaggle-on-custom-dataset/data
+def detect_Custom(img):
+    model='passport_mrz' # Naming Convention for yolov7 See output file of https://www.kaggle.com/code/owaiskhan9654/training-yolov7-on-kaggle-on-custom-dataset/data
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', nargs='+', type=str, default=model+".pt", help='model.pt path(s)')
     parser.add_argument('--source', type=str, default='Inference/', help='source') 
     parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
-    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
+    parser.add_argument('--conf-thres', type=float, default=0.45, help='object confidence threshold')
     parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
     parser.add_argument('--view-img', action='store_true', help='display results')
@@ -84,8 +90,8 @@ def detect_Custom(img,model):
 
         # Inference
         t1 = time_synchronized()
-        pred = model(img, augment=opt.augment)[0]
-
+        with torch.no_grad():  
+            pred = model(img, augment=opt.augment)[0]
 
         pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
         t2 = time_synchronized()
@@ -93,7 +99,7 @@ def detect_Custom(img,model):
         
         if classify:
             pred = apply_classifier(pred, modelc, img, im0s)
-    
+
         for i, det in enumerate(pred): 
             if webcam: 
                 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
@@ -124,6 +130,15 @@ def detect_Custom(img,model):
                     if save_img or view_img:
                         label = f'{names[int(cls)]} {conf:.2f}'
                         plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
+                        if(cls == 1):
+                          x1 = int(xyxy[0].item())
+                          y1 = int(xyxy[1].item())
+                          x2 = int(xyxy[2].item())
+                          y2 = int(xyxy[3].item())
+                          orig_img = im0
+                          crop_img = im0[y1:y2, x1:x2]
+                          cv2.imwrite('MRZ_1.png', crop_img)
+                        
             if view_img:
                 cv2.imshow(str(p), im0)
                 cv2.waitKey(1)  
@@ -146,16 +161,25 @@ def detect_Custom(img,model):
                         vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                     vid_writer.write(im0)
 
+    output_text = 'This is not a valid Passport'
+    text = pytesseract.image_to_string(Image.open('MRZ_1.png'))
+    text = text.replace(" ", "")
+    text=text[22:28]
+    today = date.today()
+    s = today.strftime('%Y%m%d')[2:]
+    if(text > s):
+      output_text = 'This is a Valid Passport'
+   
     if save_txt or save_img:
         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
 
     print(f'Done. ({time.time() - t0:.3f}s)')
 
-    return Image.fromarray(im0[:,:,::-1])
+    return [Image.fromarray(im0[:,:,::-1]), output_text]
 
    
-
-Custom_description="<center>Custom Training Performed on Kaggle <a href='https://www.kaggle.com/code/owaiskhan9654/training-yolov7-on-kaggle-on-custom-dataset/notebook' style='text-decoration: underline' target='_blank'>Link</a> </center><br> <center>Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors </center> <br> <b>1st</b> class is for Person Detected<br><b>2nd</b> class is for Car Detected"
+output = gr.Textbox(label="Validation",elem_id="opbox")
+Custom_description="<center>Custom Training Performed on Colab <a href='https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/train-yolov7-object-detection-on-custom-data.ipynb?authuser=2#scrollTo=1iqOPKjr22mL' style='text-decoration: underline' target='_blank'>Link</a> </center><br> <center>Model trained with test dataset of 'aadhar-card', 'credit-card','prescription' and 'passport' </center>"
 
 Footer = (
     "<center>Model Trained by: Owais Ahmad Data Scientist at <b> Thoucentric </b> <a href=\"https://www.linkedin.com/in/owaiskhan9654/\">Visit Profile</a> <br></center>"
@@ -169,7 +193,7 @@ Footer = (
 
 examples1=[["Image1.jpeg", "Yolo_v7_Custom_trained_By_Owais"],["Image2.jpeg", "Yolo_v7_Custom_trained_By_Owais"],["Image3.jpeg", "Yolo_v7_Custom_trained_By_Owais",],["Image4.jpeg", "Yolo_v7_Custom_trained_By_Owais"],["Image5.jpeg", "Yolo_v7_Custom_trained_By_Owais"],["Image6.jpeg", "Yolo_v7_Custom_trained_By_Owais"],["horses.jpeg", "yolov7"],["horses.jpeg", "yolov7-e6"]]
 
-Top_Title="<center>Yolov7 🚀 Custom Trained by <a href='https://www.linkedin.com/in/owaiskhan9654/' style='text-decoration: underline' target='_blank'>Owais Ahmad </center></a>🚗Car and 👦Person Detection Class"
+Top_Title="<center>Intelligent Image to Text - IIT </center></a>"
 
 css = ".output-image, .input-image, .image-preview {height: 300px !important}"
-gr.Interface(detect_Custom,[gr.Image(type="pil"),gr.Dropdown(default="Yolo_v7_Custom_trained_By_Owais",choices=["Yolo_v7_Custom_trained_By_Owais","yolov7","yolov7-e6"])],gr.Image(type="pil"),css=css,title=Top_Title,examples=examples1,description=Custom_description,article=Footer,cache_examples=False).launch()
+gr.Interface(detect_Custom,gr.Image(type="pil"),[gr.Image(type="pil"),output],css=css,title=Top_Title,examples=examples1,description=Custom_description,article=Footer,cache_examples=False).launch()