1import argparse
2
3import matplotlib.pyplot as plt
4import numpy as np
5from ai_edge_litert.compiled_model import CompiledModel
6from ai_edge_litert.hardware_accelerator import HardwareAccelerator
7from huggingface_hub import hf_hub_download
8from PIL import Image
9
10VOC_CATEGORIES = [
11 "__background__",
12 "aeroplane",
13 "bicycle",
14 "bird",
15 "boat",
16 "bottle",
17 "bus",
18 "car",
19 "cat",
20 "chair",
21 "cow",
22 "diningtable",
23 "dog",
24 "horse",
25 "motorbike",
26 "person",
27 "pottedplant",
28 "sheep",
29 "sofa",
30 "train",
31 "tvmonitor",
32]
33
34
35def _download_model_from_hf(repo_id: str, filename: str | None = None) -> str:
36 if filename:
37 return hf_hub_download(repo_id=repo_id, filename=filename)
38 candidates = ("fcn_resnet50_nchw.tflite", "fcn_resnet50.tflite")
39 last_error: Exception | None = None
40 for name in candidates:
41 try:
42 return hf_hub_download(repo_id=repo_id, filename=name)
43 except Exception as err: # pylint: disable=broad-except
44 last_error = err
45 raise FileNotFoundError(
46 f"Could not find expected model in {repo_id}: {', '.join(candidates)}"
47 ) from last_error
48
49
50def _load_cpu_model(model_path: str) -> CompiledModel:
51 return CompiledModel.from_file(model_path, hardware_accel=HardwareAccelerator.CPU)
52
53
54def _infer_nchw_input_hw(model: CompiledModel) -> tuple[int, int]:
55 req = model.get_input_buffer_requirements(0, 0)
56 dims = req.get("dimensions") or req.get("shape") or req.get("dims")
57 if not dims:
58 return 520, 520
59 dims = [int(v) for v in dims]
60 if len(dims) == 4 and dims[1] == 3:
61 return dims[2], dims[3]
62 if len(dims) == 3 and dims[0] == 3:
63 return dims[1], dims[2]
64 return 520, 520
65
66
67def _preprocess_nchw(image: Image.Image, input_h: int, input_w: int) -> np.ndarray:
68 image = image.convert("RGB")
69 w, h = image.size
70 short = 520
71 if w < h:
72 image = image.resize((short, int(round(h * short / w))), Image.BILINEAR)
73 else:
74 image = image.resize((int(round(w * short / h)), short), Image.BILINEAR)
75 image = image.resize((input_w, input_h), Image.BILINEAR)
76 x = np.asarray(image, dtype=np.float32) / 255.0
77 x = (x - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array(
78 [0.229, 0.224, 0.225], dtype=np.float32
79 )
80 return np.transpose(x, (2, 0, 1))
81
82
83def _run_logits_hwc(model: CompiledModel, nchw_input: np.ndarray, num_classes: int) -> np.ndarray:
84 inp = model.create_input_buffers(0)
85 out = model.create_output_buffers(0)
86 inp[0].write(nchw_input)
87 model.run_by_index(0, inp, out)
88
89 req = model.get_output_buffer_requirements(0, 0)
90 y = out[0].read(req["buffer_size"] // np.dtype(np.float32).itemsize, np.float32).reshape(-1)
91 h, w = nchw_input.shape[1], nchw_input.shape[2]
92 if y.size != num_classes * h * w:
93 raise ValueError(f"Unexpected output size {y.size}; expected {num_classes * h * w}")
94 chw = y.reshape(num_classes, h, w)
95 return np.transpose(chw, (1, 2, 0))
96
97
98def _softmax_last_axis(logits: np.ndarray) -> np.ndarray:
99 logits = logits.astype(np.float32, copy=False)
100 max_logits = np.max(logits, axis=-1, keepdims=True)
101 exps = np.exp(logits - max_logits)
102 return exps / np.sum(exps, axis=-1, keepdims=True)
103
104
105def _normalize_mask(mask_prob: np.ndarray) -> np.ndarray:
106 mask = np.clip(mask_prob.astype(np.float32), 0.0, 1.0)
107 lo = float(np.percentile(mask, 5.0))
108 hi = float(np.percentile(mask, 99.0))
109 if hi <= lo:
110 hi = float(mask.max())
111 lo = float(mask.min())
112 if hi <= lo:
113 return np.zeros_like(mask, dtype=np.float32)
114 return np.clip((mask - lo) / (hi - lo), 0.0, 1.0)
115
116
117def _build_overlay_rgb(base_rgb: np.ndarray, mask_prob: np.ndarray) -> np.ndarray:
118 """Returns RGB overlay image in uint8."""
119 mask = _normalize_mask(mask_prob)
120 heat = np.zeros_like(base_rgb, dtype=np.float32)
121 heat[..., 0] = 255.0
122 heat[..., 1] = 180.0 * mask
123 alpha = 0.75 * (np.clip((mask - 0.35) / 0.65, 0.0, 1.0) ** 1.5)[..., None]
124 out = (1.0 - alpha) * base_rgb.astype(np.float32) + alpha * heat
125 return np.clip(out, 0, 255).astype(np.uint8)
126
127
128def main() -> int:
129 ap = argparse.ArgumentParser()
130 ap.add_argument("--image", required=True)
131 ap.add_argument("--repo_id", default="litert-community/fcn_resnet50")
132 ap.add_argument(
133 "--model_file",
134 default=None,
135 help="Optional model filename in repo. If omitted, common names are tried.",
136 )
137 ap.add_argument("--class_name", default="dog", choices=VOC_CATEGORIES)
138 ap.add_argument(
139 "--save_figure",
140 default=None,
141 help="Optional output PNG path for the matplotlib figure.",
142 )
143 ap.add_argument(
144 "--no_show",
145 action="store_true",
146 help="Do not open interactive window (use with --save_figure).",
147 )
148 args = ap.parse_args()
149
150 model_path = _download_model_from_hf(args.repo_id, filename=args.model_file)
151 model = _load_cpu_model(model_path)
152
153 image = Image.open(args.image).convert("RGB")
154 input_h, input_w = _infer_nchw_input_hw(model)
155 x_nchw = _preprocess_nchw(image, input_h, input_w)
156 logits_hwc = _run_logits_hwc(model, x_nchw, len(VOC_CATEGORIES))
157 probs_hwc = _softmax_last_axis(logits_hwc)
158 class_index = VOC_CATEGORIES.index(args.class_name)
159 mask_prob = probs_hwc[..., class_index]
160
161 # Resize mask to original image for display.
162 mask_img = Image.fromarray((np.clip(mask_prob, 0.0, 1.0) * 255.0).astype(np.uint8), mode="L")
163 mask_img = mask_img.resize(image.size, Image.BILINEAR)
164 mask_prob_resized = np.asarray(mask_img, dtype=np.float32) / 255.0
165
166 base_rgb = np.asarray(image, dtype=np.uint8)
167 overlay_rgb = _build_overlay_rgb(base_rgb, mask_prob_resized)
168
169 fig, axs = plt.subplots(1, 3, figsize=(15, 5))
170 axs[0].imshow(base_rgb)
171 axs[0].set_title("Input")
172 axs[0].axis("off")
173
174 im = axs[1].imshow(mask_prob_resized, cmap="magma", vmin=0.0, vmax=1.0)
175 axs[1].set_title(f"Class Prob: {args.class_name}")
176 axs[1].axis("off")
177 fig.colorbar(im, ax=axs[1], fraction=0.046, pad=0.04)
178
179 axs[2].imshow(overlay_rgb)
180 axs[2].set_title("Overlay")
181 axs[2].axis("off")
182
183 fig.suptitle(f"Model: {model_path}\nClass: {args.class_name}")
184 fig.tight_layout()
185
186 if args.save_figure:
187 fig.savefig(args.save_figure, dpi=180, bbox_inches="tight")
188 print(f"Saved figure to: {args.save_figure}")
189
190 if not args.no_show:
191 plt.show()
192 else:
193 plt.close(fig)
194
195 return 0
196
197
198if __name__ == "__main__":
199 raise SystemExit(main())
