Spaces:
Sleeping
Sleeping
File size: 4,642 Bytes
5f8adf2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 |
import gradio as gr
from typing import Any, Dict, List
from src.registry import get_model_display_names, get_model
APP_TITLE = "Machine Learning CS 6140 Project: Pet Recognizer"
TOP_K_DEFAULT = 5
DARK_CSS = """
body {
background-color: #0f172a !important;
}
.gradio-container {
background-color: #0f172a !important;
color: #e5e7eb !important;
}
h1, h2, h3, h4, p, li, label {
color: #e5e7eb !important;
}
a {
color: #60a5fa !important;
}
.gr-box {
background-color: #020617 !important;
border-radius: 10px;
}
.gr-button {
background-color: #1e293b !important;
color: #e5e7eb !important;
}
.gr-button:hover {
background-color: #334155 !important;
}
"""
# -----------------------------
# Helpers
# -----------------------------
def format_topk_for_table(top_k: List[Dict[str, Any]]) -> List[List[Any]]:
rows = []
for rank, entry in enumerate(top_k, start=1):
class_name = entry.get("class_name", f"id={entry.get('class_id', '?')}")
prob = entry.get("probability", 0.0)
rows.append([rank, class_name, round(float(prob) * 100.0, 2)])
return rows
def run_inference(model_id: str, image) -> Dict[str, Any]:
if image is None:
return {
"main_text": "Please upload an image first.",
"topk_table": [],
}
model = get_model(model_id)
result = model.predict(image, top_k=TOP_K_DEFAULT)
class_name = result.get("class_name", "Unknown")
class_id = result.get("class_id", "N/A")
top_k = result.get("top_k", [])
main_text = (
f"**Predicted Class:** {class_name} \n"
f"**Class ID:** {class_id}"
)
return {
"main_text": main_text,
"topk_table": format_topk_for_table(top_k),
}
# -----------------------------
# UI
# -----------------------------
def build_demo() -> gr.Blocks:
model_display_names = get_model_display_names()
name_to_id = {v: k for k, v in model_display_names.items()}
default_display_name = next(iter(name_to_id.keys()))
with gr.Blocks(css=DARK_CSS) as demo:
# Title
gr.Markdown(
f"""
# {APP_TITLE}
This project demonstrates **pet breed recognition** using the
**Oxford-IIIT Pet Dataset**, comparing **classical machine learning models**
(Logistic Regression, SVM) with **deep feature-based models**
(Pretrained ResNet18).
**Dataset & Supported Breeds**
The models are trained on **37 cat and dog breeds** from the Oxford-IIIT Pet Dataset.
https://www.robots.ox.ac.uk/~vgg/data/pets/
"""
)
# Instructions
gr.Markdown(
"""
## Instructions
1. **Upload** a clear, close-up image of a **cat or dog** belonging to one of the supported breeds
2. **Select a model** to run the recognition:
- **LR / SVM** → Expected to perform poorly on raw pixel inputs
- **ResNet-based models** → Use pretrained deep visual features and produce much better results
3. Click **Run Identification** to view the **Top-5 predictions**
"""
)
with gr.Row():
# Left column
with gr.Column(scale=1):
gr.Markdown("### Select Model & Upload Image")
model_dropdown = gr.Dropdown(
choices=list(name_to_id.keys()),
value=default_display_name,
label="Select Model",
)
image_input = gr.Image(
type="pil",
label="Upload your pet image (JPEG / PNG)",
)
run_button = gr.Button("Run Identification")
# Right column
with gr.Column(scale=1):
gr.Markdown("### Model Prediction")
main_output = gr.Markdown(
value="Prediction will appear here.",
)
topk_output = gr.Dataframe(
headers=["Rank", "Class Name", "Probability (%)"],
datatype=["number", "str", "number"],
column_count=3,
label=f"Top-{TOP_K_DEFAULT} Predictions",
)
# Button wiring
def _gradio_infer(selected_display_name, img):
model_id = name_to_id[selected_display_name]
result = run_inference(model_id, img)
return result["main_text"], result["topk_table"]
run_button.click(
fn=_gradio_infer,
inputs=[model_dropdown, image_input],
outputs=[main_output, topk_output],
)
return demo
if __name__ == "__main__":
demo = build_demo()
demo.launch()
|