StampDetection / object_detector.py
HimankJ's picture
Added model file and code
d652175 verified
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