Spaces:
Runtime error
Runtime error
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 |