Create demo.py

demo.py

@@ -0,0 +1,173 @@
+import gradio as gr
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+from transformers import AutoModelForCausalLM
+import matplotlib
+
+matplotlib.use("Agg")  # Use Agg backend for non-interactive plotting
+
+
+os.environ["HF_TOKEN"] = os.environ.get("TOKEN_FROM_SECRET") or True
+model = AutoModelForCausalLM.from_pretrained(
+    "vikhyatk/moondream-next",
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
+    device_map={"": "cuda"},
+    revision="69420e0c6596863b4f0059e365fadc5cb388e8fd"
+)
+
+
+def visualize_gaze_multi(face_boxes, gaze_points, image=None, show_plot=True):
+    """Visualization function with reduced whitespace"""
+    # Calculate figure size based on image aspect ratio
+    if image is not None:
+        height, width = image.shape[:2]
+        aspect_ratio = width / height
+        fig_height = 6  # Base height
+        fig_width = fig_height * aspect_ratio
+    else:
+        width, height = 800, 600
+        fig_width, fig_height = 10, 8
+
+    # Create figure with tight layout
+    fig = plt.figure(figsize=(fig_width, fig_height))
+    ax = fig.add_subplot(111)
+
+    if image is not None:
+        ax.imshow(image)
+    else:
+        ax.set_facecolor("#1a1a1a")
+        fig.patch.set_facecolor("#1a1a1a")
+
+    colors = plt.cm.rainbow(np.linspace(0, 1, len(face_boxes)))
+
+    for face_box, gaze_point, color in zip(face_boxes, gaze_points, colors):
+        hex_color = "#{:02x}{:02x}{:02x}".format(
+            int(color[0] * 255), int(color[1] * 255), int(color[2] * 255)
+        )
+
+        x, y, width_box, height_box = face_box
+        gaze_x, gaze_y = gaze_point
+
+        face_center_x = x + width_box / 2
+        face_center_y = y + height_box / 2
+
+        face_rect = plt.Rectangle(
+            (x, y), width_box, height_box, fill=False, color=hex_color, linewidth=2
+        )
+        ax.add_patch(face_rect)
+
+        points = 50
+        alphas = np.linspace(0.8, 0, points)
+
+        x_points = np.linspace(face_center_x, gaze_x, points)
+        y_points = np.linspace(face_center_y, gaze_y, points)
+
+        for i in range(points - 1):
+            ax.plot(
+                [x_points[i], x_points[i + 1]],
+                [y_points[i], y_points[i + 1]],
+                color=hex_color,
+                alpha=alphas[i],
+                linewidth=4,
+            )
+
+        ax.scatter(gaze_x, gaze_y, color=hex_color, s=100, zorder=5)
+        ax.scatter(gaze_x, gaze_y, color="white", s=50, zorder=6)
+
+    # Set plot limits and remove axes
+    ax.set_xlim(0, width)
+    ax.set_ylim(height, 0)
+    ax.set_aspect("equal")
+    ax.set_xticks([])
+    ax.set_yticks([])
+
+    # Remove padding around the plot
+    plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
+
+    return fig
+
+@spaces.GPU(duration=15)
+def process_image(input_image):
+    try:
+        # Convert to PIL Image if needed
+        if isinstance(input_image, np.ndarray):
+            pil_image = Image.fromarray(input_image)
+        else:
+            pil_image = input_image
+
+        # Get image encoding
+        enc_image = model.encode_image(pil_image)
+
+        # Detect faces
+        faces = model.detect(enc_image, "face")["objects"]
+
+        if not faces:
+            return None, "No faces detected in the image."
+
+        # Process each face
+        face_boxes = []
+        gaze_points = []
+
+        for face in faces:
+            face_center = (
+                (face["x_min"] + face["x_max"]) / 2,
+                (face["y_min"] + face["y_max"]) / 2,
+            )
+            gaze = model.detect_gaze(enc_image, face_center)
+
+            if gaze is None:
+                continue
+
+            face_box = (
+                face["x_min"] * pil_image.width,
+                face["y_min"] * pil_image.height,
+                (face["x_max"] - face["x_min"]) * pil_image.width,
+                (face["y_max"] - face["y_min"]) * pil_image.height,
+            )
+
+            gaze_point = (
+                gaze["x"] * pil_image.width,
+                gaze["y"] * pil_image.height,
+            )
+
+            face_boxes.append(face_box)
+            gaze_points.append(gaze_point)
+
+        # Create visualization
+        image_array = np.array(pil_image)
+        fig = visualize_gaze_multi(
+            face_boxes, gaze_points, image=image_array, show_plot=False
+        )
+
+        return fig, f"Detected {len(faces)} faces."
+
+    except Exception as e:
+        return None, f"Error processing image: {str(e)}"
+
+
+with gr.Blocks(title="Moondream Gaze Detection") as app:
+    gr.Markdown("# 🌔 Moondream Gaze Detection")
+    gr.Markdown("Upload an image to detect faces and visualize their gaze directions.")
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Input Image", type="pil")
+
+        with gr.Column():
+            output_text = gr.Textbox(label="Status")
+            output_plot = gr.Plot(label="Visualization")
+
+    input_image.change(
+        fn=process_image, inputs=[input_image], outputs=[output_plot, output_text]
+    )
+
+    gr.Examples(
+        examples=["gaze_test.jpg", "gaze_test2.jpg", "gaze_test3.jpg"],
+        inputs=input_image,
+    )
+
+if __name__ == "__main__":
+    app.launch()