Spaces:
Build error
Build error
File size: 1,944 Bytes
f2f4624 | 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 | from typing import List, Dict, Any
import io
import numpy as np
from PIL import Image
import requests
import tensorflow as tf
# 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_bytes: bytes) -> np.ndarray:
# Mirror TS: ensure JPEG-like decode and resize 224x224, keep 0..255 range
image = Image.open(io.BytesIO(image_bytes))
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(self, inputs: bytes) -> List[Dict[str, Any]]:
x = _preprocess(inputs)
preds = self.model.predict(x)
if isinstance(preds, (list, tuple)):
preds = preds[0]
probs = np.asarray(preds).squeeze().tolist()
# Top-1 output following TS behavior
idx = int(np.argmax(probs))
return [
{"label": LABELS[idx], "score": float(probs[idx])},
]
def load_model(model_dir: str = ".") -> PreTrainedModel:
# HF Inference API convention: a top-level load entrypoint
return PreTrainedModel(model_path=f"{model_dir}/model/model_resnet50.keras")
|