Update app/app_savta.py

app/app_savta.py

@@ -1,106 +1,152 @@
-import torch
 import os
+from pathlib import Path
+
+import torch
 from fastai.vision.all import *
 import gradio as gr
 
-############### HF ###########################
-
-HF_TOKEN = os.getenv('HF_TOKEN')
+######################
+# Hugging Face flags #
+######################
 
+HF_TOKEN = os.getenv("HF_TOKEN")
+# `HuggingFaceDatasetSaver` is still available in Gradio ≥ 5
 hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "savtadepth-flags-V2")
 
-############## DVC ################################
+############
+#   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"):
+if Path(".dvc").is_dir():
     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			
+    if os.system(f"dvc pull {PROD_MODEL_PATH} {TRAIN_PATH} {TEST_PATH}") != 0:
+        raise SystemExit("dvc pull failed")
+    # remove DVC metadata to avoid accidental reuse in the Space
+    os.system("rm -rf .dvc")
 
-############## Inference ##############################
+#######################
+#   Data & Learner    #
+#######################
 
 class ImageImageDataLoaders(DataLoaders):
-    """Basic wrapper around several `DataLoader`s with factory methods for Image to Image problems"""
+    """Wrapper to create DataLoaders for image→image tasks."""
+
     @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 ################################
+    def from_label_func(
+        cls,
+        path: Path,
+        filenames,
+        label_func,
+        valid_pct: float = 0.2,
+        seed: int | None = None,
+        item_transforms=None,
+        batch_transforms=None,
+        **kwargs,
+    ):
+        dblock = 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,
+        )
+        return cls.from_dblock(dblock, filenames, path=path, **kwargs)
+
+
+def get_y_fn(x: Path) -> Path:
+    """Map an RGB image path to its depth‑map counterpart."""
+    return Path(str(x).replace(".jpg", "_depth.png"))
+
+
+def create_data(data_path: Path):
+    fnames = get_files(data_path / "train", extensions=".jpg")
+    return ImageImageDataLoaders.from_label_func(
+        data_path / "train",
+        seed=42,
+        bs=4,
+        num_workers=0,
+        filenames=fnames,
+        label_func=get_y_fn,
+    )
+
+
+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")
+
+#####################
+#  Inference Logic  #
+#####################
+
+def predict_depth(input_img: PILImage) -> PILImageBW:
+    """Generate a single‑channel depth prediction from an RGB image."""
+    depth, *_ = learner.predict(input_img)
+    return PILImageBW.create(depth).convert("L")
+
+#####################
+#    Gradio UI      #
+#####################
 
 title = "SavtaDepth WebApp"
 
-description = """
-<p>
-<center>
-Savta Depth is a collaborative Open Source Data Science project for monocular depth estimation - Turn 2d photos into 3d photos. To test the model and code please check out the link bellow.
-<img src="https://huggingface.co/spaces/kingabzpro/savtadepth/resolve/main/examples/cover.png" alt="logo" width="250"/>
-</center>
-</p>
-"""
-article = "<p style='text-align: center'><a href='https://dagshub.com/OperationSavta/SavtaDepth' target='_blank'>SavtaDepth Project from OperationSavta</a></p><p style='text-align: center'><a href='https://colab.research.google.com/drive/1XU4DgQ217_hUMU1dllppeQNw3pTRlHy1?usp=sharing' target='_blank'>Google Colab Demo</a></p></center><center><img src='https://visitor-badge.glitch.me/badge?page_id=kingabzpro/savtadepth' alt='visitor badge'></center></p>"
+description = (
+    """
+    <p style="text-align:center;">
+        Savta Depth is a collaborative Open‑Source project for monocular depth estimation – turn 2‑D photos into 3‑D. 🏞️<br>
+        Try the model below or explore the resources.
+        <br><img src="https://huggingface.co/spaces/kingabzpro/savtadepth/resolve/main/examples/cover.png" width="250"/>
+    </p>
+    """
+)
+
+article = (
+    """
+    <p style='text-align:center'>
+        <a href='https://dagshub.com/OperationSavta/SavtaDepth' target='_blank'>Project on DAGsHub</a> •
+        <a href='https://colab.research.google.com/drive/1XU4DgQ217_hUMU1dllppeQNw3pTRlHy1?usp=sharing' target='_blank'>Google Colab Demo</a>
+        <br/>
+        <img src='https://visitor-badge.glitch.me/badge?page_id=kingabzpro/savtadepth' alt='visitor badge'/>
+    </p>
+    """
+)
 
 examples = [
     ["examples/00008.jpg"],
     ["examples/00045.jpg"],
 ]
-favicon = "examples/favicon.ico"
-thumbnail = "examples/SavtaDepth.png"
-
-
-def main():
-	iface = gr.Interface(
-	                     gen, 
-	                     gr.inputs.Image(shape=(640,480),type='numpy'), 
-	                     "image", 
-	                     title = title, 
-	                     flagging_options=["incorrect", "worst","ambiguous"], 
-	                     allow_flagging = "manual",
-	                     flagging_callback=hf_writer,
-	                     description = description, 
-	                     article = article, 
-	                     examples = examples, 
-	                     theme ="peach",
-	                     allow_screenshot=True
-	                     )
-
-	iface.launch(enable_queue=True)
-# enable_queue=True,auth=("admin", "pass1234")
-
-if __name__ == '__main__':
-	main()
+
+# Modern Gradio components (v5+)
+input_component = gr.Image(shape=(640, 480), image_mode="RGB", label="Input RGB")
+output_component = gr.Image(image_mode="L", label="Predicted Depth")
+
+with gr.Blocks(title=title) as demo:
+    gr.Markdown(description)
+    gr.Markdown(article)
+
+    gr.Interface(
+        fn=predict_depth,
+        inputs=input_component,
+        outputs=output_component,
+        allow_flagging="manual",
+        flagging_options=["incorrect", "worst", "ambiguous"],
+        flagging_callback=hf_writer,
+        examples=examples,
+        cache_examples=False,  # use live inference for examples
+        theme=gr.themes.Soft(),
+    )
+
+if __name__ == "__main__":
+    demo.queue(concurrency_count=3).launch()