Update app.py

app.py

@@ -1,76 +1,86 @@
-import gradio as gr
-import numpy as np
-from PIL import Image
-import tensorflow as tf
-from typing import List, Dict, Any
-import io
-
-# Labels must mirror src/classification-model/index.ts
-LABELS: List[str] = [
-    "battery",
-    "biological",
-    "brown-glass",
-    "cardboard",
-    "clothes",
-    "green-glass",
-    "metal",
-    "paper",
-    "plastic",
-    "shoes",
-    "trash",
-    "white-glass",
-]
-
-
-def _load_image_to_rgb(image: Image.Image) -> np.ndarray:
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    return np.asarray(image)
-
-
-def _resize_224(img_rgb: np.ndarray) -> np.ndarray:
-    im = Image.fromarray(img_rgb)
-    im = im.resize((224, 224), Image.NEAREST)
-    return np.asarray(im)
-
-
-def _preprocess(image: Image.Image) -> np.ndarray:
-    rgb = _load_image_to_rgb(image)
-    rgb224 = _resize_224(rgb)
-    # shape [1,224,224,3], float32 in 0..255
-    arr = rgb224.astype("float32")
-    return np.expand_dims(arr, axis=0)
-
-
-class PreTrainedModel:
-    def __init__(self, model_path: str = "model/model_resnet50.keras") -> None:
-        self.model = tf.keras.models.load_model(model_path)
-
-    def predict_image(self, image: Image.Image) -> Dict[str, float]:
-        x = _preprocess(image)
-        preds = self.model.predict(x)
-        if isinstance(preds, (list, tuple)):
-            preds = preds[0]
-        probs = np.asarray(preds).squeeze().tolist()
-
-        return {label: score for label, score in zip(LABELS, probs)}
-
-
-model = PreTrainedModel()
-
-
-def predict(image):
-    predictions = model.predict_image(image)
-    return predictions
-
-
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=3),
-    title="Waste Classification",
-    description="Upload an image of waste to classify it.",
-)
-
-if __name__ == "__main__":
-    iface.launch()
+import gradio as gr
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+from typing import List, Dict, Any
+import io
+
+# Labels must mirror src/classification-model/index.ts
+LABELS: List[str] = [
+    "battery",
+    "biological",
+    "brown-glass",
+    "cardboard",
+    "clothes",
+    "green-glass",
+    "metal",
+    "paper",
+    "plastic",
+    "shoes",
+    "trash",
+    "white-glass",
+]
+
+
+def _load_image_to_rgb(image: Image.Image) -> np.ndarray:
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    return np.asarray(image)
+
+
+def _resize_224(img_rgb: np.ndarray) -> np.ndarray:
+    im = Image.fromarray(img_rgb)
+    im = im.resize((224, 224), Image.NEAREST)
+    return np.asarray(im)
+
+
+def _preprocess(image: Image.Image) -> np.ndarray:
+    rgb = _load_image_to_rgb(image)
+    rgb224 = _resize_224(rgb)
+    # shape [1,224,224,3], float32 in 0..255
+    arr = rgb224.astype("float32")
+    return np.expand_dims(arr, axis=0)
+
+
+class PreTrainedModel:
+    def __init__(self, model_path: str = "model/model_resnet50.keras") -> None:
+        self.model = tf.keras.models.load_model(model_path)
+
+    def predict_image(self, image: Image.Image) -> Dict[str, float]:
+        x = _preprocess(image)
+        preds = self.model.predict(x)
+        if isinstance(preds, (list, tuple)):
+            preds = preds[0]
+        probs = np.asarray(preds).squeeze().tolist()
+
+        return {label: score for label, score in zip(LABELS, probs)}
+
+
+model = PreTrainedModel()
+
+
+def predict(image):
+    predictions = model.predict_image(image)
+
+    probs_percent = {label: round(p * 100, 2)
+                     for label, p in predictions.items()}
+
+    max_label = max(probs_percent, key=probs_percent.get)
+
+    return {
+        "label": max_label,
+        "percentage": probs_percent[max_label],
+        "probabilities": probs_percent,
+    }
+
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.JSON(),
+    title="Waste Classification",
+    description="Upload an image of waste to classify it.",
+)
+
+if __name__ == "__main__":
+    iface.launch()