import torch import os from fastai.vision.all import * import gradio as gr ############### HF ########################### HF_TOKEN = os.getenv('HF_TOKEN') hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "savtadepth-flags-V2") ############## DVC ################################ PROD_MODEL_PATH = "src/models" TRAIN_PATH = "src/data/processed/train/bathroom" TEST_PATH = "src/data/processed/test/bathroom" if os.path.isdir(".dvc"): print("Running DVC") # os.system("dvc config cache.type copy") # os.system("dvc config core.no_scm true") if os.system(f"dvc pull {PROD_MODEL_PATH} {TRAIN_PATH } {TEST_PATH }") != 0: exit("dvc pull failed") os.system("rm -r .dvc") # .apt/usr/lib/dvc ############## Inference ############################## class ImageImageDataLoaders(DataLoaders): """Basic wrapper around several `DataLoader`s with factory methods for Image to Image problems""" @classmethod @delegates(DataLoaders.from_dblock) def from_label_func(cls, path, filenames, label_func, valid_pct=0.2, seed=None, item_transforms=None, batch_transforms=None, **kwargs): """Create from list of `fnames` in `path`s with `label_func`.""" datablock = DataBlock(blocks=(ImageBlock(cls=PILImage), ImageBlock(cls=PILImageBW)), get_y=label_func, splitter=RandomSplitter(valid_pct, seed=seed), item_tfms=item_transforms, batch_tfms=batch_transforms) res = cls.from_dblock(datablock, filenames, path=path, **kwargs) return res def get_y_fn(x): y = str(x.absolute()).replace('.jpg', '_depth.png') y = Path(y) return y def create_data(data_path): fnames = get_files(data_path/'train', extensions='.jpg') data = ImageImageDataLoaders.from_label_func(data_path/'train', seed=42, bs=4, num_workers=0, filenames=fnames, label_func=get_y_fn) return data data = create_data(Path('src/data/processed')) learner = unet_learner(data,resnet34, metrics=rmse, wd=1e-2, n_out=3, loss_func=MSELossFlat(), path='src/') learner.load('model') def gen(input_img): return PILImageBW.create((learner.predict(input_img))[0]).convert('L') ################### Gradio Web APP ################################ title = "SavtaDepth WebApp" description = """
SavtaDepth Project from OperationSavta