Initial commit for Coral AI Demo

.idea/workspace.xml

@@ -42,7 +42,7 @@
       <option name="number" value="Default" />
       <option name="presentableId" value="Default" />
       <updated>1733405525071</updated>
-      <workItem from="1733405526612" duration="1184000" />
+      <workItem from="1733405526612" duration="1615000" />
     </task>
     <servers />
   </component>

README.md

@@ -1,12 +1,13 @@
 ---
-title: Coral Ai Demo
+title: Coral AI Demo
 emoji: 🐠
 colorFrom: blue
 colorTo: green
-sdk: static
-pinned: false
+sdk: gradio
+app_file: app.py
+pinned: true
 license: mit
-short_description: Demo Spaces for our Coral AI Models
+short_description: Interact with Coral AI models for image analysis and segmentation.
 ---
 
 # Coral AI Demo

app.py

@@ -1,73 +1,83 @@
 import gradio as gr
 import torch
-import requests
 import os
+from ultralytics import YOLO
 
-# Hugging Face models repository URL
-HF_MODELS_REPO = "https://huggingface.co/reefsupport/coral-ai/resolve/main/models"
-
-# Available models
-MODEL_VERSIONS = [
+# Coral AI model files hosted in the Hugging Face model repository
+model_names = [
     "yolov8_xlarge_latest.pt",
     "yolov8_xlarge_v1.pt",
-    "yolov8_xlarge_v2.pt"
+    "yolov8_xlarge_v2.pt",
 ]
 
+# Set the initial model
+current_model_name = "yolov8_xlarge_latest.pt"
+model_dir = "models"
+os.makedirs(model_dir, exist_ok=True)
 
-# Download the selected model if not present locally
-def download_model(version):
-    model_path = f"models/{version}"
+# Download models if not already present locally
+HF_MODELS_REPO = "https://huggingface.co/reefsupport/coral-ai/resolve/main/models"
+for model_name in model_names:
+    model_path = os.path.join(model_dir, model_name)
     if not os.path.exists(model_path):
-        os.makedirs("models", exist_ok=True)
-        url = f"{HF_MODELS_REPO}/{version}"
-        print(f"Downloading {version} from {url}...")
-        response = requests.get(url, stream=True)
-        with open(model_path, "wb") as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                if chunk:
-                    f.write(chunk)
-        print(f"Downloaded {version} to {model_path}")
-    return model_path
+        print(f"Downloading {model_name}...")
+        model_url = f"{HF_MODELS_REPO}/{model_name}"
+        torch.hub.download_url_to_file(model_url, model_path)
 
+# Load the initial model
+model = YOLO(os.path.join(model_dir, current_model_name))
 
-# Load the model
-def load_model(version):
-    model_path = download_model(version)
-    model = torch.load(model_path, map_location=torch.device("cpu"))
-    model.eval()
-    return model
 
+def coral_ai_inference(image: str, model_name: str):
+    """
+    Coral AI inference function
+    Args:
+        image: Input image filepath
+        model_name: Name of the model
+    Returns:
+        Rendered image
+    """
+    global model
+    global current_model_name
+    if model_name != current_model_name:
+        model = YOLO(os.path.join(model_dir, model_name))
+        current_model_name = model_name
 
-# Perform inference
-def infer(version, image):
-    try:
-        model = load_model(version)
-        # Placeholder logic: replace with actual model inference
-        return f"Model '{version}' processed the uploaded image successfully."
-    except Exception as e:
-        return f"Error during inference: {str(e)}"
+    # Perform inference
+    results = model.predict(image, return_outputs=True)
 
+    # Render results (placeholder logic)
+    rendered_image = results[0].plot()  # Visualization of predictions
+    return rendered_image
 
-# Set up Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Coral AI Demo")
-    gr.Markdown("Select a model version and upload an image for inference.")
 
-    with gr.Row():
-        version_dropdown = gr.Dropdown(choices=MODEL_VERSIONS, label="Select Model Version")
-        image_input = gr.Image(type="filepath", label="Upload Image")
-
-    with gr.Row():
-        result_output = gr.Textbox(label="Inference Result")
+# Define Gradio inputs and outputs
+inputs = [
+    gr.Image(type="filepath", label="Input Image"),
+    gr.Dropdown(
+        model_names,
+        value=current_model_name,
+        label="Model Type",
+    ),
+]
 
-    with gr.Row():
-        run_button = gr.Button("Run Model")
+outputs = gr.Image(type="filepath", label="Output Image")
+title = "Coral AI YOLOv8 Segmentation Demo"
 
-    # Link components
-    run_button.click(
-        fn=infer,
-        inputs=[version_dropdown, image_input],
-        outputs=result_output
-    )
+examples = [
+    ["examples/coral_image1.jpg", "yolov8_xlarge_latest.pt"],
+    ["examples/coral_image2.jpg", "yolov8_xlarge_latest.pt"],
+    ["examples/coral_image3.jpg", "yolov8_xlarge_latest.pt"],
+]
 
-demo.launch()
+# Create and launch the Gradio interface
+demo_app = gr.Interface(
+    fn=coral_ai_inference,
+    inputs=inputs,
+    outputs=outputs,
+    title=title,
+    examples=examples,
+    cache_examples=True,
+    theme="default",
+)
+demo_app.queue().launch(debug=True)