Minor changes

app.py

@@ -23,12 +23,8 @@ import postprocess
 
 
 ocr_instance = PaddleOCR(use_angle_cls=False, lang='en', use_gpu=True)
-
-torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
-detection_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/detection_wts.pt', force_reload=True)
-
-torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
-structure_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/structure_wts.pt', force_reload=True)
+detection_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/detection_wts.pt', force_reload=True, skip_validation=True, trust_repo=True)
+structure_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/structure_wts.pt', force_reload=True, skip_validation=True, trust_repo=True)
 
 imgsz = 640
 
@@ -87,13 +83,18 @@ def crop_image(pil_img, detection_result, padding=30):
         w = result[2]
         h = result[3]
 
-        x1 = max(0, int((min_x - w / 2) * width) - padding)
-        y1 = max(0, int((min_y - h / 2) * height) - padding)
-        x2 = min(width, int((min_x + w / 2) * width) + padding)
-        y2 = min(height, int((min_y + h / 2) * height) + padding)
+        x1 = int((min_x - w / 2) * width)
+        y1 = int((min_y - h / 2) * height)
+        x2 = int((min_x + w / 2) * width)
+        y2 = int((min_y + h / 2) * height)
         # print(x1, y1, x2, y2)
 
-        crop_image = image[y1:y2, x1:x2, :]
+        x1_pad = max(0, x1 - padding)
+        y1_pad = max(0, y1 - padding)
+        x2_pad = min(width, x2 + padding)
+        y2_pad = min(height, y2 + padding)
+
+        crop_image = image[y1_pad:y2_pad, x1_pad:x2_pad, :]
         crop_image = cv_to_PIL(crop_image)
         if class_id == 1:  # table rotated
             crop_image = crop_image.rotate(270, expand=True)
@@ -101,6 +102,7 @@ def crop_image(pil_img, detection_result, padding=30):
         crop_images.append(crop_image)
 
         cv2.rectangle(image, (x1, y1), (x2, y2), color=(0, 0, 255))
+        cv2.putText(image, f'{score:.2f}', (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(255, 0, 0))
 
     return crop_images, cv_to_PIL(image)