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