modify app

README.md

@@ -7,5 +7,3 @@ sdk: gradio
 app_file: app.py
 pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -16,14 +16,12 @@ from focusondepth.model_definition import FocusOnDepth
 AutoConfig.register("focusondepth", FocusOnDepthConfig)
 AutoModel.register(FocusOnDepthConfig, FocusOnDepth)
 
-original_image_cache = {}
 transform = transforms.Compose([
     transforms.Resize((384, 384)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
 ])
 model = AutoModel.from_pretrained('ybelkada/focusondepth', trust_remote_code=True)
-# model.load_state_dict(torch.load('./focusondepth/FocusOnDepth_vit_base_patch16_384.p', map_location=torch.device('cpu'))['model_state_dict'])
 
 @torch.no_grad()
 def inference(input_image):
@@ -42,6 +40,20 @@ def inference(input_image):
 
     return [depth.resize(original_size, resample=Image.BICUBIC), segmentation.resize(original_size, resample=Image.NEAREST)]
 
+description = """
+<center>
+Can a single model predict both segmentation and depth estimation? At least, if the segmentation is constrained for a single class, the answer is yes! <br>
+In this project, we use a DPT model to predict the depth and the segmentation mask of the class human, of an image. This model could be potentially used for an autofocus application where you would need the segmentation mask of the humans on the picture, as well as the depth estimation of the scene<br>
+</center>
+"""
+title="""
+FocusOnDepth - A single DPT encoder for Dense Prediction Tasks
+"""
+css = """
+"""
+article = "<div style='text-align: center;'><a href='https://github.com/isl-org/DPT' target='_blank'>Original Paper</a> | <a href='https://github.com/antocad/FocusOnDepth' target='_blank'>Extended Version</a></div>"
+
+
 iface = gr.Interface(
     fn=inference, 
     inputs=gr.inputs.Image(label="Input Image"), 
@@ -49,5 +61,9 @@ iface = gr.Interface(
         gr.outputs.Image(label="Depth Map:"),
         gr.outputs.Image(label="Segmentation Map:"),
     ],
+    description=description,
+    title=title,
+    css=css,
+    article=article,
 )
 iface.launch()