Update app.py

app.py

@@ -1,58 +1,52 @@
-from doctest import Example
 import gradio as gr
 from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 import torch
 import numpy as np
 from PIL import Image
-import open3d as o3d
-from pathlib import Path
-import os
 
+# Load model và feature extractor
 feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
 model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
+model.eval()
 
+def process_image(image):
+    # Đảm bảo ảnh là RGB
+    if image.mode != "RGB":
+        image = image.convert("RGB")
 
-def process_image(image_path):
-    image_path = Path(image_path)
-    image_raw = Image.open(image_path)
-    image = image_raw.resize(
-        (1920, int(1920 * image_raw.size[1] / image_raw.size[0])),
-        Image.Resampling.LANCZOS,
-    )
-    # prepare image for the model
+    # Encode
     encoding = feature_extractor(image, return_tensors="pt")
-
-    # forward pass
+    
+    # Dự đoán depth
     with torch.no_grad():
         outputs = model(**encoding)
         predicted_depth = outputs.predicted_depth
-
-    # interpolate to original size
+    
+    # Resize về kích thước ảnh gốc
     prediction = torch.nn.functional.interpolate(
         predicted_depth.unsqueeze(1),
         size=image.size[::-1],
         mode="bicubic",
-        align_corners=False,
+        align_corners=False
     ).squeeze()
+    
+    # Chuẩn hóa và chuyển về ảnh uint8
     output = prediction.cpu().numpy()
-    depth_image = (output * 255 / np.max(output)).astype("uint8")
-    img = Image.fromarray(depth_image)
-    return [img]
+    output = (output - np.min(output)) / (np.max(output) - np.min(output))  # normalize
+    formatted = (output * 255).astype("uint8")
+    depth_img = Image.fromarray(formatted)
 
+    return depth_img
 
-
-title = "Demo: zero-shot depth estimation with DPT + 3D Point Cloud"
-description = "This demo is a variation from the original <a href='https://huggingface.co/spaces/nielsr/dpt-depth-estimation' target='_blank'>DPT Demo</a>. It uses the DPT model to predict the depth of an image and then uses 3D Point Cloud to create a 3D object."
-examples = [["examples/" + img] for img in os.listdir("examples/")]
+title = "Demo: zero-shot depth estimation with DPT"
+description = "Demo for Intel's DPT, a Dense Prediction Transformer for state-of-the-art dense prediction tasks such as semantic segmentation and depth estimation."
 
 iface = gr.Interface(
     fn=process_image,
-    inputs=[gr.inputs.Image(type="filepath", label="Input Image")],
-    outputs=[gr.outputs.Image(type="filepath", label="Predicted Depth")],
+    inputs=gr.inputs.Image(type="pil", label="Input Image"),
+    outputs=gr.outputs.Image(type="pil", label="Predicted Depth"),
     title=title,
     description=description,
-    examples=examples,
-    allow_flagging="never",
-    #cache_examples=False,
 )
-iface.launch()
+
+iface.launch(debug=True)