initial commit

AUTHORS.rst

@@ -0,0 +1,13 @@
+=======
+Credits
+=======
+
+Development Lead
+----------------
+
+* Full name of the author <Email of the author>
+
+Contributors
+------------
+
+None yet. Why not be the first?

README.md

@@ -1,12 +1,22 @@
 ---
-title: Table To Csv Pipeline
-emoji: 🏃
-colorFrom: purple
-colorTo: blue
+title: Table Extraction
+emoji: 🚀
+colorFrom: indigo
+colorTo: purple
 sdk: streamlit
-sdk_version: 1.28.1
+sdk_version: 1.21.0
 app_file: app.py
 pinned: false
 ---
+# huggingface-space
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Imported from https://huggingface.co/spaces/jurgendn/table-extraction with some adjustment.
+
+
+Current pipeline:
+
+Table detection: https://huggingface.co/microsoft/table-transformer-detection
+
+Table recognition: https://huggingface.co/microsoft/table-transformer-structure-recognition
+
+OCR: https://github.com/pbcquoc/vietocr

app.py

@@ -0,0 +1,570 @@
+import asyncio
+import string
+from collections import Counter
+from itertools import count, tee
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import streamlit as st
+import torch
+from PIL import Image
+from transformers import DetrImageProcessor, TableTransformerForObjectDetection
+from vietocr.tool.config import Cfg
+from vietocr.tool.predictor import Predictor
+
+st.set_option('deprecation.showPyplotGlobalUse', False)
+st.set_page_config(layout='wide')
+st.title("Table Detection and Table Structure Recognition")
+st.write(
+ "Implemented by MSFT team: https://github.com/microsoft/table-transformer")
+
+# config = Cfg.load_config_from_name('vgg_transformer')
+config = Cfg.load_config_from_name('vgg_seq2seq')
+config['cnn']['pretrained'] = False
+config['device'] = 'cpu'
+config['predictor']['beamsearch'] = False
+detector = Predictor(config)
+
+table_detection_model = TableTransformerForObjectDetection.from_pretrained(
+ "microsoft/table-transformer-detection")
+
+table_recognition_model = TableTransformerForObjectDetection.from_pretrained(
+ "microsoft/table-transformer-structure-recognition")
+
+
+def PIL_to_cv(pil_img):
+ return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+
+
+def cv_to_PIL(cv_img):
+ return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
+
+
+async def pytess(cell_pil_img, threshold: float = 0.5):
+ text, prob = detector.predict(cell_pil_img, return_prob=True)
+ if prob < threshold:
+ return ""
+ return text.strip()
+
+
+def sharpen_image(pil_img):
+
+ img = PIL_to_cv(pil_img)
+ sharpen_kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
+
+ sharpen = cv2.filter2D(img, -1, sharpen_kernel)
+ pil_img = cv_to_PIL(sharpen)
+ return pil_img
+
+
+def uniquify(seq, suffs=count(1)):
+ """Make all the items unique by adding a suffix (1, 2, etc).
+ Credit: https://stackoverflow.com/questions/30650474/python-rename-duplicates-in-list-with-progressive-numbers-without-sorting-list
+ `seq` is mutable sequence of strings.
+ `suffs` is an optional alternative suffix iterable.
+ """
+ not_unique = [k for k, v in Counter(seq).items() if v > 1]
+
+ suff_gens = dict(zip(not_unique, tee(suffs, len(not_unique))))
+ for idx, s in enumerate(seq):
+ try:
+ suffix = str(next(suff_gens[s]))
+ except KeyError:
+ continue
+ else:
+ seq[idx] += suffix
+
+ return seq
+
+
+def binarizeBlur_image(pil_img):
+ image = PIL_to_cv(pil_img)
+ thresh = cv2.threshold(image, 150, 255, cv2.THRESH_BINARY_INV)[1]
+
+ result = cv2.GaussianBlur(thresh, (5, 5), 0)
+ result = 255 - result
+ return cv_to_PIL(result)
+
+
+def td_postprocess(pil_img):
+ '''
+ Removes gray background from tables
+ '''
+ img = PIL_to_cv(pil_img)
+
+ hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+ mask = cv2.inRange(hsv, (0, 0, 100),
+ (255, 5, 255)) # (0, 0, 100), (255, 5, 255)
+ nzmask = cv2.inRange(hsv, (0, 0, 5),
+ (255, 255, 255)) # (0, 0, 5), (255, 255, 255))
+ nzmask = cv2.erode(nzmask, np.ones((3, 3))) # (3,3)
+ mask = mask & nzmask
+
+ new_img = img.copy()
+ new_img[np.where(mask)] = 255
+
+ return cv_to_PIL(new_img)
+
+
+# def super_res(pil_img):
+# # requires opencv-contrib-python installed without the opencv-python
+# sr = dnn_superres.DnnSuperResImpl_create()
+# image = PIL_to_cv(pil_img)
+# model_path = "./LapSRN_x8.pb"
+# model_name = model_path.split('/')[1].split('_')[0].lower()
+# model_scale = int(model_path.split('/')[1].split('_')[1].split('.')[0][1])
+
+# sr.readModel(model_path)
+# sr.setModel(model_name, model_scale)
+# final_img = sr.upsample(image)
+# final_img = cv_to_PIL(final_img)
+
+# return final_img
+
+
+def table_detector(image, THRESHOLD_PROBA):
+ '''
+ Table detection using DEtect-object TRansformer pre-trained on 1 million tables
+
+ '''
+
+ feature_extractor = DetrImageProcessor(do_resize=True,
+ size=800,
+ max_size=800)
+ encoding = feature_extractor(image, return_tensors="pt")
+
+ with torch.no_grad():
+ outputs = table_detection_model(**encoding)
+
+ probas = outputs.logits.softmax(-1)[0, :, :-1]
+ keep = probas.max(-1).values > THRESHOLD_PROBA
+
+ target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)
+ postprocessed_outputs = feature_extractor.post_process(
+ outputs, target_sizes)
+ bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]
+
+ return (probas[keep], bboxes_scaled)
+
+
+def table_struct_recog(image, THRESHOLD_PROBA):
+ '''
+ Table structure recognition using DEtect-object TRansformer pre-trained on 1 million tables
+ '''
+
+ feature_extractor = DetrImageProcessor(do_resize=True,
+ size=1000,
+ max_size=1000)
+ encoding = feature_extractor(image, return_tensors="pt")
+
+ with torch.no_grad():
+ outputs = table_recognition_model(**encoding)
+
+ probas = outputs.logits.softmax(-1)[0, :, :-1]
+ keep = probas.max(-1).values > THRESHOLD_PROBA
+
+ target_sizes = torch.tensor(image.size[::-1]).unsqueeze(0)
+ postprocessed_outputs = feature_extractor.post_process(
+ outputs, target_sizes)
+ bboxes_scaled = postprocessed_outputs[0]['boxes'][keep]
+
+ return (probas[keep], bboxes_scaled)
+
+
+class TableExtractionPipeline():
+
+ colors = ["red", "blue", "green", "yellow", "orange", "violet"]
+
+ # colors = ["red", "blue", "green", "red", "red", "red"]
+
+ def add_padding(self,
+ pil_img,
+ top,
+ right,
+ bottom,
+ left,
+ color=(255, 255, 255)):
+ '''
+ Image padding as part of TSR pre-processing to prevent missing table edges
+ '''
+ width, height = pil_img.size
+ new_width = width + right + left
+ new_height = height + top + bottom
+ result = Image.new(pil_img.mode, (new_width, new_height), color)
+ result.paste(pil_img, (left, top))
+ return result
+
+ def plot_results_detection(self, c1, model, pil_img, prob, boxes,
+ delta_xmin, delta_ymin, delta_xmax, delta_ymax):
+ '''
+ crop_tables and plot_results_detection must have same co-ord shifts because 1 only plots the other one updates co-ordinates 
+ '''
+ # st.write('img_obj')
+ # st.write(pil_img)
+ plt.imshow(pil_img)
+ ax = plt.gca()
+
+ for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
+ cl = p.argmax()
+ xmin, ymin, xmax, ymax = xmin - delta_xmin, ymin - delta_ymin, xmax + delta_xmax, ymax + delta_ymax
+ ax.add_patch(
+ plt.Rectangle((xmin, ymin),
+ xmax - xmin,
+ ymax - ymin,
+ fill=False,
+ color='red',
+ linewidth=3))
+ text = f'{model.config.id2label[cl.item()]}: {p[cl]:0.2f}'
+ ax.text(xmin - 20,
+ ymin - 50,
+ text,
+ fontsize=10,
+ bbox=dict(facecolor='yellow', alpha=0.5))
+ plt.axis('off')
+ c1.pyplot()
+
+ def crop_tables(self, pil_img, prob, boxes, delta_xmin, delta_ymin,
+ delta_xmax, delta_ymax):
+ '''
+ crop_tables and plot_results_detection must have same co-ord shifts because 1 only plots the other one updates co-ordinates 
+ '''
+ cropped_img_list = []
+
+ for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
+
+ xmin, ymin, xmax, ymax = xmin - delta_xmin, ymin - delta_ymin, xmax + delta_xmax, ymax + delta_ymax
+ cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+ cropped_img_list.append(cropped_img)
+
+ return cropped_img_list
+
+ def generate_structure(self, c2, model, pil_img, prob, boxes,
+ expand_rowcol_bbox_top, expand_rowcol_bbox_bottom):
+ '''
+ Co-ordinates are adjusted here by 3 'pixels'
+ To plot table pillow image and the TSR bounding boxes on the table
+ '''
+ # st.write('img_obj')
+ # st.write(pil_img)
+ plt.figure(figsize=(32, 20))
+ plt.imshow(pil_img)
+ ax = plt.gca()
+ rows = {}
+ cols = {}
+ idx = 0
+
+ for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
+
+ xmin, ymin, xmax, ymax = xmin, ymin, xmax, ymax
+ cl = p.argmax()
+ class_text = model.config.id2label[cl.item()]
+ text = f'{class_text}: {p[cl]:0.2f}'
+ # or (class_text == 'table column')
+ if (class_text
+ == 'table row') or (class_text
+ == 'table projected row header') or (
+ class_text == 'table column'):
+ ax.add_patch(
+ plt.Rectangle((xmin, ymin),
+ xmax - xmin,
+ ymax - ymin,
+ fill=False,
+ color=self.colors[cl.item()],
+ linewidth=2))
+ ax.text(xmin - 10,
+ ymin - 10,
+ text,
+ fontsize=5,
+ bbox=dict(facecolor='yellow', alpha=0.5))
+
+ if class_text == 'table row':
+ rows['table row.' +
+ str(idx)] = (xmin, ymin - expand_rowcol_bbox_top, xmax,
+ ymax + expand_rowcol_bbox_bottom)
+ if class_text == 'table column':
+ cols['table column.' +
+ str(idx)] = (xmin, ymin - expand_rowcol_bbox_top, xmax,
+ ymax + expand_rowcol_bbox_bottom)
+
+ idx += 1
+
+ plt.axis('on')
+ c2.pyplot()
+ return rows, cols
+
+ def sort_table_featuresv2(self, rows: dict, cols: dict):
+ # Sometimes the header and first row overlap, and we need the header bbox not to have first row's bbox inside the headers bbox
+ rows_ = {
+ table_feature: (xmin, ymin, xmax, ymax)
+ for table_feature, (
+ xmin, ymin, xmax,
+ ymax) in sorted(rows.items(), key=lambda tup: tup[1][1])
+ }
+ cols_ = {
+ table_feature: (xmin, ymin, xmax, ymax)
+ for table_feature, (
+ xmin, ymin, xmax,
+ ymax) in sorted(cols.items(), key=lambda tup: tup[1][0])
+ }
+
+ return rows_, cols_
+
+ def individual_table_featuresv2(self, pil_img, rows: dict, cols: dict):
+
+ for k, v in rows.items():
+ xmin, ymin, xmax, ymax = v
+ cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+ rows[k] = xmin, ymin, xmax, ymax, cropped_img
+
+ for k, v in cols.items():
+ xmin, ymin, xmax, ymax = v
+ cropped_img = pil_img.crop((xmin, ymin, xmax, ymax))
+ cols[k] = xmin, ymin, xmax, ymax, cropped_img
+
+ return rows, cols
+
+ def object_to_cellsv2(self, master_row: dict, cols: dict,
+ expand_rowcol_bbox_top, expand_rowcol_bbox_bottom,
+ padd_left):
+ '''Removes redundant bbox for rows&columns and divides each row into cells from columns
+ Args:
+
+ Returns:
+
+ 
+ '''
+ cells_img = {}
+ header_idx = 0
+ row_idx = 0
+ previous_xmax_col = 0
+ new_cols = {}
+ new_master_row = {}
+ previous_ymin_row = 0
+ new_cols = cols
+ new_master_row = master_row
+ ## Below 2 for loops remove redundant bounding boxes ###
+ # for k_col, v_col in cols.items():
+ # xmin_col, _, xmax_col, _, col_img = v_col
+ # if (np.isclose(previous_xmax_col, xmax_col, atol=5)) or (xmin_col >= xmax_col):
+ # print('Found a column with double bbox')
+ # continue
+ # previous_xmax_col = xmax_col
+ # new_cols[k_col] = v_col
+
+ # for k_row, v_row in master_row.items():
+ # _, ymin_row, _, ymax_row, row_img = v_row
+ # if (np.isclose(previous_ymin_row, ymin_row, atol=5)) or (ymin_row >= ymax_row):
+ # print('Found a row with double bbox')
+ # continue
+ # previous_ymin_row = ymin_row
+ # new_master_row[k_row] = v_row
+ ######################################################
+ for k_row, v_row in new_master_row.items():
+
+ _, _, _, _, row_img = v_row
+ xmax, ymax = row_img.size
+ xa, ya, xb, yb = 0, 0, 0, ymax
+ row_img_list = []
+ # plt.imshow(row_img)
+ # st.pyplot()
+ for idx, kv in enumerate(new_cols.items()):
+ k_col, v_col = kv
+ xmin_col, _, xmax_col, _, col_img = v_col
+ xmin_col, xmax_col = xmin_col - padd_left - 10, xmax_col - padd_left
+ xa = xmin_col
+ xb = xmax_col
+ if idx == 0:
+ xa = 0
+ if idx == len(new_cols) - 1:
+ xb = xmax
+ xa, ya, xb, yb = xa, ya, xb, yb
+
+ row_img_cropped = row_img.crop((xa, ya, xb, yb))
+ row_img_list.append(row_img_cropped)
+
+ cells_img[k_row + '.' + str(row_idx)] = row_img_list
+ row_idx += 1
+
+ return cells_img, len(new_cols), len(new_master_row) - 1
+
+ def clean_dataframe(self, df):
+ '''
+ Remove irrelevant symbols that appear with tesseractOCR
+ '''
+ # df.columns = [col.replace('|', '') for col in df.columns]
+
+ for col in df.columns:
+
+ df[col] = df[col].str.replace("'", '', regex=True)
+ df[col] = df[col].str.replace('"', '', regex=True)
+ df[col] = df[col].str.replace(']', '', regex=True)
+ df[col] = df[col].str.replace('[', '', regex=True)
+ df[col] = df[col].str.replace('{', '', regex=True)
+ df[col] = df[col].str.replace('}', '', regex=True)
+ return df
+
+ @st.cache
+ def convert_df(self, df):
+ return df.to_csv().encode('utf-8')
+
+ def create_dataframe(self, c3, cell_ocr_res: list, max_cols: int,
+ max_rows: int):
+ '''Create dataframe using list of cell values of the table, also checks for valid header of dataframe
+ Args:
+ cell_ocr_res: list of strings, each element representing a cell in a table
+ max_cols, max_rows: number of columns and rows
+ Returns:
+ dataframe : final dataframe after all pre-processing 
+ '''
+
+ headers = cell_ocr_res[:max_cols]
+ new_headers = uniquify(headers,
+ (f' {x!s}' for x in string.ascii_lowercase))
+ counter = 0
+
+ cells_list = cell_ocr_res[max_cols:]
+ df = pd.DataFrame("", index=range(0, max_rows), columns=new_headers)
+
+ cell_idx = 0
+ for nrows in range(max_rows):
+ for ncols in range(max_cols):
+ df.iat[nrows, ncols] = str(cells_list[cell_idx])
+ cell_idx += 1
+
+ ## To check if there are duplicate headers if result of uniquify+col == col
+ ## This check removes headers when all headers are empty or if median of header word count is less than 6
+ for x, col in zip(string.ascii_lowercase, new_headers):
+ if f' {x!s}' == col:
+ counter += 1
+ header_char_count = [len(col) for col in new_headers]
+
+ # if (counter == len(new_headers)) or (statistics.median(header_char_count) < 6):
+ # st.write('woooot')
+ # df.columns = uniquify(df.iloc[0], (f' {x!s}' for x in string.ascii_lowercase))
+ # df = df.iloc[1:,:]
+
+ df = self.clean_dataframe(df)
+
+ c3.dataframe(df)
+ csv = self.convert_df(df)
+ c3.download_button("Download table",
+ csv,
+ "file.csv",
+ "text/csv",
+ key='download-csv-' + df.iloc[0, 0])
+
+ return df
+
+ async def start_process(self, image_path: str, TD_THRESHOLD, TSR_THRESHOLD,
+ OCR_THRESHOLD, padd_top, padd_left, padd_bottom,
+ padd_right, delta_xmin, delta_ymin, delta_xmax,
+ delta_ymax, expand_rowcol_bbox_top,
+ expand_rowcol_bbox_bottom):
+ '''
+ Initiates process of generating pandas dataframes from raw pdf-page images
+
+ '''
+ image = Image.open(image_path).convert("RGB")
+ probas, bboxes_scaled = table_detector(image,
+ THRESHOLD_PROBA=TD_THRESHOLD)
+
+ if bboxes_scaled.nelement() == 0:
+ st.write('No table found in the pdf-page image')
+ return ''
+
+ # try:
+ # st.write('Document: '+image_path.split('/')[-1])
+ c1, c2, c3 = st.columns((1, 1, 1))
+
+ self.plot_results_detection(c1, table_detection_model, image, probas,
+ bboxes_scaled, delta_xmin, delta_ymin,
+ delta_xmax, delta_ymax)
+ cropped_img_list = self.crop_tables(image, probas, bboxes_scaled,
+ delta_xmin, delta_ymin, delta_xmax,
+ delta_ymax)
+
+ for idx, unpadded_table in enumerate(cropped_img_list):
+
+ table = self.add_padding(unpadded_table, padd_top, padd_right,
+ padd_bottom, padd_left)
+ # table = super_res(table)
+ # table = binarizeBlur_image(table)
+ # table = sharpen_image(table) # Test sharpen image next
+ # table = td_postprocess(table)
+
+ # table.save("result"+str(idx)+".png")
+
+ probas, bboxes_scaled = table_struct_recog(
+ table, THRESHOLD_PROBA=TSR_THRESHOLD)
+ rows, cols = self.generate_structure(c2, table_recognition_model,
+ table, probas, bboxes_scaled,
+ expand_rowcol_bbox_top,
+ expand_rowcol_bbox_bottom)
+ # st.write(len(rows), len(cols))
+ rows, cols = self.sort_table_featuresv2(rows, cols)
+ master_row, cols = self.individual_table_featuresv2(
+ table, rows, cols)
+
+ cells_img, max_cols, max_rows = self.object_to_cellsv2(
+ master_row, cols, expand_rowcol_bbox_top,
+ expand_rowcol_bbox_bottom, padd_left)
+
+ sequential_cell_img_list = []
+ for k, img_list in cells_img.items():
+ for img in img_list:
+ # img = super_res(img)
+ # img = sharpen_image(img) # Test sharpen image next
+ # img = binarizeBlur_image(img)
+ # img = self.add_padding(img, 10,10,10,10)
+ # plt.imshow(img)
+ # c3.pyplot()
+ sequential_cell_img_list.append(
+ pytess(cell_pil_img=img, threshold=OCR_THRESHOLD))
+
+ cell_ocr_res = await asyncio.gather(*sequential_cell_img_list)
+
+ self.create_dataframe(c3, cell_ocr_res, max_cols, max_rows)
+ st.write(
+ 'Errors in OCR is due to either quality of the image or performance of the OCR'
+ )
+ # except:
+ # st.write('Either incorrectly identified table or no table, to debug remove try/except')
+ # break
+ # break
+
+
+if __name__ == "__main__":
+
+ img_name = st.file_uploader("Upload an image with table(s)")
+ st1, st2, st3 = st.columns((1, 1, 1))
+ TD_th = st1.slider('Table detection threshold', 0.0, 1.0, 0.8)
+ TSR_th = st2.slider('Table structure recognition threshold', 0.0, 1.0, 0.7)
+ OCR_th = st3.slider("Text Probs Threshold", 0.0, 1.0, 0.5)
+
+ st1, st2, st3, st4 = st.columns((1, 1, 1, 1))
+
+ padd_top = st1.slider('Padding top', 0, 200, 90)
+ padd_left = st2.slider('Padding left', 0, 200, 40)
+ padd_right = st3.slider('Padding right', 0, 200, 40)
+ padd_bottom = st4.slider('Padding bottom', 0, 200, 90)
+
+ te = TableExtractionPipeline()
+ # for img in image_list:
+ if img_name is not None:
+ asyncio.run(
+ te.start_process(img_name,
+ TD_THRESHOLD=TD_th,
+ TSR_THRESHOLD=TSR_th,
+ OCR_THRESHOLD=OCR_th,
+ padd_top=padd_top,
+ padd_left=padd_left,
+ padd_bottom=padd_bottom,
+ padd_right=padd_right,
+ delta_xmin=10, # add offset to the left of the table
+ delta_ymin=3, # add offset to the bottom of the table
+ delta_xmax=10, # add offset to the right of the table
+ delta_ymax=3, # add offset to the top of the table
+ expand_rowcol_bbox_top=0,
+ expand_rowcol_bbox_bottom=0))

components/callbacks.py

@@ -0,0 +1,4 @@
+# Define callbacks here
+from pytorch_lightning.callbacks import EarlyStopping
+
+early_stopping = EarlyStopping(monitor="loss", min_delta=0, patience=3)

components/data_module.py

@@ -0,0 +1,81 @@
+from typing import Callable, List, Optional, Union
+
+import torch
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, Dataset
+
+
+class SampleDataset(Dataset):
+
+ def __init__(self,
+ x: Union[List, torch.Tensor],
+ y: Union[List, torch.Tensor],
+ transforms: Optional[Callable] = None) -> None:
+ super(SampleDataset, self).__init__()
+ self.x = x
+ self.y = y
+
+ if transforms is None:
+ # Replace None with some default transforms
+ # If image, could be an Resize and ToTensor
+ self.transforms = lambda x: x
+ else:
+ self.transforms = transforms
+
+ def __len__(self):
+ return len(self.x)
+
+ def __getitem__(self, index: int):
+ x = self.x[index]
+ y = self.y[index]
+
+ x = self.transforms(x)
+ return x, y
+
+
+class SampleDataModule(LightningDataModule):
+
+ def __init__(self,
+ x: Union[List, torch.Tensor],
+ y: Union[List, torch.Tensor],
+ transforms: Optional[Callable] = None,
+ val_ratio: float = 0,
+ batch_size: int = 32) -> None:
+ super(SampleDataModule, self).__init__()
+ assert 0 <= val_ratio < 1
+ assert isinstance(batch_size, int)
+ self.x = x
+ self.y = y
+
+ self.transforms = transforms
+ self.val_ratio = val_ratio
+ self.batch_size = batch_size
+
+ self.setup()
+ self.prepare_data()
+
+ def setup(self, stage: Optional[str] = None) -> None:
+ pass
+
+ def prepare_data(self) -> None:
+ n_samples: int = len(self.x)
+ train_size: int = n_samples - int(n_samples * self.val_ratio)
+
+ self.train_dataset = SampleDataset(x=self.x[:train_size],
+ y=self.y[:train_size],
+ transforms=self.transforms)
+ if train_size < n_samples:
+ self.val_dataset = SampleDataset(x=self.x[train_size:],
+ y=self.y[train_size:],
+ transforms=self.transforms)
+ else:
+ self.val_dataset = SampleDataset(x=self.x[-self.batch_size:],
+ y=self.y[-self.batch_size:],
+ transforms=self.transforms)
+
+ def train_dataloader(self) -> DataLoader:
+ return DataLoader(dataset=self.train_dataset,
+ batch_size=self.batch_size)
+
+ def val_dataloader(self) -> DataLoader:
+ return DataLoader(dataset=self.val_dataset, batch_size=self.batch_size)

config.py

@@ -0,0 +1,3 @@
+from dynaconf import Dynaconf
+
+CFG = Dynaconf(envvar_prefix="DYNACONF", settings_files=["config/config.yaml"])

docker-compose.yml

@@ -0,0 +1,23 @@
+version: "3.7"
+
+services:
+ model_name:
+ build:
+ context: .
+ dockerfile: .docker/Dockerfile
+ container_name: model_name
+ ports:
+ - "8996:8996"
+ env_file:
+ - ./.env
+ volumes:
+ - ./data:/home/working/data:ro
+
+ # This part is used to enable GPU support
+ deploy:
+ resources:
+ reservations:
+ devices:
+ - driver: nvidia
+ count: 1
+ capabilities: [ gpu ]

models/base_model/classification.py

@@ -0,0 +1,63 @@
+from abc import abstractmethod
+from typing import Any, Dict, List
+
+import torch
+from pytorch_lightning import LightningModule
+from torch import Tensor
+
+
+class LightningClassification(LightningModule):
+
+ @abstractmethod
+ def __init__(self, *args, **kwargs) -> None:
+ super(LightningClassification, self).__init__(*args, **kwargs)
+ self.train_batch_output: List[Dict] = []
+ self.validation_batch_output: List[Dict] = []
+ self.log_value_list: List[str] = ['loss', 'f1', 'precision', 'recall']
+
+ @abstractmethod
+ def forward(self, *args, **kwargs) -> Any:
+ pass
+
+ @abstractmethod
+ def configure_optimizers(self):
+ pass
+
+ @abstractmethod
+ def loss(self, input: Tensor, target: Tensor, **kwargs) -> Tensor:
+ pass
+
+ @abstractmethod
+ def training_step(self, batch, batch_idx):
+ pass
+
+ def __average(self, key: str, outputs: List[Dict]) -> Tensor:
+ target_arr = torch.Tensor([val[key] for val in outputs]).float()
+ return target_arr.mean()
+
+ @torch.no_grad()
+ def on_train_epoch_start(self) -> None:
+ self.train_batch_output = []
+
+ @torch.no_grad()
+ def on_train_epoch_end(self) -> None:
+ for key in self.log_value_list:
+ val = self.__average(key=key, outputs=self.train_batch_output)
+ log_name = f"training/{key}"
+ self.log(name=log_name, value=val)
+
+ @abstractmethod
+ @torch.no_grad()
+ def validation_step(self, batch, batch_idx):
+ pass
+
+ @torch.no_grad()
+ def on_validation_epoch_start(self) -> None:
+ self.validation_batch_output = []
+
+ @torch.no_grad()
+ def on_validation_epoch_end(self) -> None:
+ for key in self.log_value_list:
+ val = self.__average(key=key, outputs=self.validation_batch_output)
+ log_name = f"val/{key}"
+ self.log(name=log_name, value=val)

models/base_model/regression.py

@@ -0,0 +1,55 @@
+from abc import abstractmethod
+from typing import Any, Dict, List
+
+import torch
+from pytorch_lightning import LightningModule
+from torch import Tensor
+
+
+class LightningRegression(LightningModule):
+
+ @abstractmethod
+ def __init__(self, *args, **kwargs) -> None:
+ super(LightningRegression, self).__init__(*args, **kwargs)
+ self.train_step_output: List[Dict] = []
+ self.validation_step_output: List[Dict] = []
+ self.log_value_list: List[str] = ['loss', 'mse', 'mape']
+
+ @abstractmethod
+ def forward(self, *args, **kwargs) -> Any:
+ pass
+
+ @abstractmethod
+ def configure_optimizers(self):
+ pass
+
+ @abstractmethod
+ def loss(self, input: Tensor, output: Tensor, **kwargs):
+ return 0
+
+ @abstractmethod
+ def training_step(self, batch, batch_idx):
+ pass
+
+ def __average(self, key: str, outputs: List[Dict]) -> Tensor:
+ target_arr = torch.Tensor([val[key] for val in outputs]).float()
+ return target_arr.mean()
+
+ @torch.no_grad()
+ def on_train_epoch_end(self) -> None:
+ for key in self.log_value_list:
+ val = self.__average(key=key, outputs=self.train_step_output)
+ log_name = f"training/{key}"
+ self.log(name=log_name, value=val)
+
+ @torch.no_grad()
+ @abstractmethod
+ def validation_step(self, batch, batch_idx):
+ pass
+
+ @torch.no_grad()
+ def validation_epoch_end(self, outputs):
+ for key in self.log_value_list:
+ val = self.__average(key=key, outputs=self.validation_step_output)
+ log_name = f"training/{key}"
+ self.log(name=log_name, value=val)

models/metrics/classification.py

@@ -0,0 +1,44 @@
+from typing import Dict
+
+import torch
+from torchmetrics import functional as FM
+
+
+def classification_metrics(
+ preds: torch.Tensor,
+ target: torch.Tensor,
+ num_classes: int,
+ average: str = 'macro',
+ task: str = 'multiclass') -> Dict[str, torch.Tensor]:
+ """
+ get_classification_metrics
+ Return some metrics evaluation the classification task
+
+ Parameters
+ ----------
+ preds : torch.Tensor
+ logits, probs
+ target : torch.Tensor
+ targets label
+
+ Returns
+ -------
+ Dict[str, torch.Tensor]
+ _description_
+ """
+ f1 = FM.f1_score(preds=preds,
+ target=target,
+ num_classes=num_classes,
+ task=task,
+ average=average)
+ recall = FM.recall(preds=preds,
+ target=target,
+ num_classes=num_classes,
+ task=task,
+ average=average)
+ precision = FM.precision(preds=preds,
+ target=target,
+ num_classes=num_classes,
+ task=task,
+ average=average)
+ return dict(f1=f1, precision=precision, recall=recall)

models/metrics/regression.py

@@ -0,0 +1,28 @@
+from typing import Dict
+
+import torch
+from torchmetrics import functional as FM
+
+
+def regression_metrics(preds: torch.Tensor,
+ target: torch.Tensor) -> Dict[str, torch.Tensor]:
+ """
+ get_classification_metrics
+ Return some metrics evaluation the classification task
+
+ Parameters
+ ----------
+ preds : torch.Tensor
+ logits, probs
+ target : torch.Tensor
+ targets label
+
+ Returns
+ -------
+ Dict[str, torch.Tensor]
+ _description_
+ """
+ mse: torch.Tensor = FM.mean_squared_error(preds=preds, target=target)
+ mape: torch.Tensor = FM.mean_absolute_percentage_error(preds=preds,
+ target=target)
+ return dict(mse=mse, mape=mape)

models/model_lit.py

@@ -0,0 +1,50 @@
+from torch import Tensor, nn, optim
+from torch.nn import functional as F
+
+from .base_model.classification import LightningClassification
+from .metrics.classification import classification_metrics
+from .modules.sample_torch_module import UselessLayer
+
+
+class UselessClassification(LightningClassification):
+
+ def __init__(self, n_classes: int, lr: float, **kwargs) -> None:
+ super(UselessClassification).__init__()
+ self.save_hyperparameters()
+ self.n_classes = n_classes
+ self.lr = lr
+ self.main = nn.Sequential(UselessLayer(), nn.GELU())
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.main(x)
+
+ def loss(self, input: Tensor, target: Tensor) -> Tensor:
+ return F.mse_loss(input=input, target=target)
+
+ def configure_optimizers(self):
+ optimizer = optim.Adam(params=self.parameters(), lr=self.lr)
+ return optimizer
+
+ def training_step(self, batch, batch_idx):
+ x, y = batch
+
+ logits = self.forward(x)
+ loss = self.loss(input=x, target=y)
+ metrics = classification_metrics(preds=logits,
+ target=y,
+ num_classes=self.n_classes)
+
+ self.train_batch_output.append({'loss': loss, **metrics})
+ return loss
+
+ def validation_step(self, batch, batch_idx):
+ x, y = batch
+
+ logits = self.forward(x)
+ loss = self.loss(input=x, target=y)
+ metrics = classification_metrics(preds=logits,
+ target=y,
+ num_classes=self.n_classes)
+
+ self.validation_batch_output.append({'loss': loss, **metrics})
+ return loss

models/modules/sample_torch_module.py

@@ -0,0 +1,12 @@
+from torch import Tensor, nn
+
+
+class UselessLayer(nn.Module):
+
+ def __init__(self) -> None:
+ super(UselessLayer, self).__init__()
+ self.seq = nn.Identity()
+
+ def forward(self, x: Tensor) -> Tensor:
+ x = self.seq(x)
+ return x

requirements.txt

@@ -0,0 +1,9 @@
+timm==0.9.2
+torch --index-url https://download.pytorch.org/whl/cpu
+torchvision --index-url https://download.pytorch.org/whl/cpu
+torchaudio --index-url https://download.pytorch.org/whl/cpu
+vietocr==0.3.11
+streamlit==1.21.0
+pandas
+transformers==4.29.1
+Pillow==9.5.0

tests/test_resource.py

@@ -0,0 +1,4 @@
+def test_cuda():
+ from torch.cuda import is_available
+ assert is_available()
+