Added model file and code

object_detector.py

@@ -0,0 +1,46 @@
+import torch
+
+class ImageObjectDetector(object):
+    
+    def __init__(self):
+        custom_path = r"object_extractor.pt"
+        self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=custom_path, force_reload=False, source = 'github')
+        self.model.conf = 0.1
+        self.classes = self.model.names
+        self.total_time = 0
+        self.total_img_processed = 0
+        
+    def detect(self, image_path : str,  save_crop_images : bool = False, save_path : str = None):
+        pred = self.model(image_path)
+        df = pred.pandas().xyxy[0]
+        croped = pred.crop(save = save_crop_images, save_dir = save_path, exist_ok=True)
+        return df, croped
+    
+class DocumentObjects(ImageObjectDetector):
+    
+    def __init__(self):
+        super().__init__()
+        self.items_to_extract = ['stamp']
+        self.keymapper = {"sign" : "signature", "checked_item" : "checkedItem", "qr_code" : "qrCode", "bar_code" : "barCode"}
+    
+    def detect_objects(self, img_path : str):
+        df, _ = self.detect(img_path, False, None)
+
+        if df.shape[0] < 0:
+            items = []
+        else:
+            items = df['name'].unique()
+        response = {}
+        for key in self.items_to_extract:
+            temp = {}
+            loc = []
+            if key in items:
+                temp['found'] = True
+                for ind in df[df['name'] == key].index:
+                    cord = [df.loc[ind, 'xmin'], df.loc[ind, 'ymin'], df.loc[ind, 'xmax'], df.loc[ind, 'ymax']]
+                    loc.append(cord)
+            else:
+                temp['found'] = False
+            temp['loc'] = loc
+            response[self.keymapper.get(key,key)] = temp
+        return response

object_extractor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1fb3ff703e2e0367770bc1cd7431004b474e9689e9eec18f9c82c0b8be1e370
+size 168401186