Spaces:
Runtime error
Runtime error
File size: 2,133 Bytes
d28ba37 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
from ultralyticsplus import YOLO, render_result
from pathlib import Path, PurePath
class Table_extraction():
def __init__(self, image, OUTPUTS_DIR):
self.model = YOLO('keremberke/yolov8m-table-extraction')
self.model.overrides['conf'] = 0.25
self.model.overrides['iou'] = 0.45
self.model.overrides['agnostic_nms'] = False
self.model.overrides['max_det'] = 1000
self.image = image
self.OUTPUTS_DIR = OUTPUTS_DIR
def get_results(self):
self.results = self.model(self.image)
render = render_result(model=self.model, image=self.image, result=self.results[0])
op_img = self.OUTPUTS_DIR + Path(self.image).name
render.save(op_img)
return op_img
# def recognize_coords(self):
# final_results = []
# result = str(self.results)
# print(result)
# result = result.split('(')
# print('Result', result)
# if '\n' not in result[1]:
# print('In if')
# coords = []
# result = result[1][2: -3].split(',')
# for i in result:
# coords.append(float(i))
# final_results.append(coords)
# return final_results
# else:
# result = result[1:][0]
# result = result.split('\n')
# j = 0
# print('Results: ', result)
# for i in result:
# coords = []
# i = i.strip()
# if j == 0:
# print(i)
# i = i[2:-2]
# print(i)
# elif j == len(result) - 1:
# i = i[1:-3]
# else:
# i = i[1:-2]
# j+=1
# print(i)
# i = i.split(',')
# print(i)
# for k in i:
# coords.append(float(k))
# final_results.append(coords)
# return final_results
# te = Table_extraction('junks\9.png', 'outputs/123')
# te.get_results()
# print(te.recognize_coords()) |