diff --git a/workspace/sprint_1_2/CODEBASE/infra/tests/test_predict_inflammation.py b/workspace/sprint_1_2/CODEBASE/infra/tests/test_predict_inflammation.py new file mode 100644 index 0000000..0a6ee3d --- /dev/null +++ b/workspace/sprint_1_2/CODEBASE/infra/tests/test_predict_inflammation.py @@ -0,0 +1,138 @@ +import json +import time +from pathlib import Path +import numpy as np +from PIL import Image +import requests + +# Server configuration +TRITON_URL = "https://dtj-tran--triton-s3-service-unified-triton-server.modal.run" +TARGET_MODEL = "workspace/sprint_1_2/CODEBASE/infra/tests/models/efficientnet_b0_ultrasound_2_class.pth" # Target a single sub-model directly +IMFLAMMATION_CLASSES = ["Không viêm", "Có viêm"] + +BASE_DIR = Path(__file__).resolve().parent + + +def load_image(path: Path) -> Image.Image: + if not path.exists(): + raise FileNotFoundError(f"Test image not found at: {path}") + return Image.open(path).convert("RGB") + + +def preprocess_224(img: Image.Image) -> np.ndarray: + """Preprocesses image to NCHW FP32 [1, 3, 224, 224] matching ResNet50 input requirements""" + img_resized = img.resize((224, 224), Image.Resampling.BILINEAR) + arr = np.asarray(img_resized).astype(np.float32) / 255.0 + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + arr = (arr - mean) / std + arr = arr.transpose(2, 0, 1) # HWC -> CHW + arr = np.expand_dims(arr, axis=0) # Add batch dim -> NCHW + return arr + + +def softmax(x: np.ndarray) -> np.ndarray: + e = np.exp(x - np.max(x)) + return e / np.sum(e) + +def build_binary_request(input_tensor: np.ndarray, model_name: str) -> tuple[bytes, dict]: + """Constructs a strict KServe v2 binary payload for a single input tensor""" + # NOTE: Check your individual model's config.pbtxt to verify if the + # expected input name is 'input_image', 'input_0', etc. + input_name = "input_image" + + inputs = [ + { + "name": input_name, + "shape": list(input_tensor.shape), + "datatype": "FP32", + "parameters": {"binary_data_size": input_tensor.nbytes}, + } + ] + + # We request the standard 'logits' output tensor from this model + outputs = [{"name": "logits"}] + + metadata = { + "model_name": model_name, + "model_version": "", + "inputs": inputs, + "outputs": outputs + } + + metadata_bytes = json.dumps(metadata).encode("utf-8") + body = metadata_bytes + input_tensor.tobytes() + + headers = { + "Inference-Header-Content-Length": str(len(metadata_bytes)), + "Content-Type": "application/octet-stream" + } + return body, headers + +def parse_binary_response(resp: requests.Response) -> np.ndarray: + """Parses Triton's binary response stream for a single output tensor""" + header_length = int(resp.headers.get("Inference-Header-Content-Length", "0")) + response_bytes = resp.content + + metadata = json.loads(response_bytes[:header_length].decode("utf-8")) + binary_data = response_bytes[header_length:] + + # Extract the first available output tensor + output_desc = metadata["outputs"][0] + shape = output_desc["shape"] + + # Unpack raw bytes back into numpy array + arr = np.frombuffer(binary_data, dtype=np.float32).reshape(shape) + return arr + +def main(): + # Define a path to a real local image to test with + image_path = "test_images/sup-up-long_positive/58e7a7ef-de3e-11ee-97e2-0a580a5f5b60_11.png" + + print(f"Targeting Individual Model: {TARGET_MODEL}") + print(f"Loading image from: {image_path}") + + try: + # 1. Prepare image + img = load_image(image_path) + input_data = preprocess_224(img) + + # 2. Build the payload + body, headers = build_binary_request(input_data, TARGET_MODEL) + + # 3. Fire request to the targeted model endpoint + url = f"{TRITON_URL}/v2/models/{TARGET_MODEL}/infer" + print("Sending request to Triton server...") + + t0 = time.time() + resp = requests.post(url, data=body, headers=headers, timeout=30) + resp.raise_for_status() + latency = time.time() - t0 + + # 4. Parse output logits + logits = parse_binary_response(resp) + logits = np.squeeze(logits) # Drop batch dimension + + # 5. Decode probabilities + probs = softmax(logits) + predicted_idx = int(np.argmax(probs)) + + print("\n" + "=" * 40) + print("🎉 SUCCESSFUL INFERENCE") + print(f"Network Roundtrip Time: {latency:.4f}s") + print("-" * 40) + print("Class Probabilities:") + for cls_name, prob in zip(IMFLAMMATION_CLASSES, probs): + print(f" {cls_name:<15}: {prob * 100:.2f}%") + print("-" * 40) + print(f"Prediction: {IMFLAMMATION_CLASSES[predicted_idx]} ({probs[predicted_idx] * 100:.2f}%)") + print("=" * 40) + + except Exception as e: + print(f"\n❌ Inference Failed: {e}") + if 'resp' in locals() and hasattr(resp, 'text'): + print(f"Server Error Message: {resp.text}") + + +if __name__ == "__main__": + main() \ No newline at end of file