| from PIL import Image | |
| import numpy as np | |
| from utils import AppUtils | |
| import cv2 | |
| from inference import start_inference | |
| class AppInference: | |
| def __init__(self): | |
| self.COLOR_MAP = {} | |
| def inference(self, input_id, img_path, mask, label): | |
| AppUtils.clear() | |
| self._input_id = input_id | |
| self._handle_preprocess(img_path, mask, label) | |
| return self._handle_model_inference() | |
| def _handle_preprocess(self, img_path, mask, label): | |
| items = self._read_files(img_path) | |
| self._items = items | |
| mask = self._save_mask(items, mask) | |
| if label != "None": | |
| self._edit_maps(items, mask, label, save=True) | |
| def _handle_model_inference(self): | |
| return start_inference() | |
| def preview(self, input_id, img_path, mask, label): | |
| AppUtils.clear() | |
| self._input_id = input_id | |
| items = self._read_files(img_path) | |
| mask = self._save_mask(items, mask) | |
| if label != "None": | |
| self._edit_maps(items, mask, label) | |
| return self.generate_colored_image(items["inst_map"]) | |
| def generate_colored_image(self, semantic_map): | |
| np.random.seed(256) | |
| if len(semantic_map.shape) == 3: | |
| semantic_map = semantic_map[:,:,1] | |
| color_image = np.zeros((semantic_map.shape[0], semantic_map.shape[1], 3), dtype=np.uint8) | |
| for row in range(semantic_map.shape[0]): | |
| for col in range(semantic_map.shape[1]): | |
| inst_id = semantic_map[row, col] | |
| if self.COLOR_MAP.get(inst_id, None) is None: | |
| self.COLOR_MAP[inst_id] = np.random.randint(256, size=(3,)) | |
| color_image[row, col, :] = self.COLOR_MAP[inst_id] | |
| return Image.fromarray(color_image) | |
| def _read_files(self, img_path): | |
| dataset = img_path.split("/")[2] | |
| items = { | |
| "img_path": img_path, | |
| "label_path": img_path.replace("images", "labels").replace("jpg", "png"), | |
| "inst_map_path": img_path.replace("images", "inst_map").replace("jpg", "png"), | |
| } | |
| for file_path in items.values(): | |
| AppUtils.copy_file(file_path, file_path.replace(dataset, "test_processed")) | |
| items["dataset"] = dataset | |
| base_img = cv2.imread(img_path) | |
| base_img = cv2.cvtColor(base_img, cv2.COLOR_BGR2RGB) | |
| base_lab = cv2.imread(items["label_path"], 0) | |
| base_inst_map = Image.open(items["inst_map_path"]) | |
| base_inst_map = np.array(base_inst_map, dtype=np.int32) | |
| items.update( | |
| { | |
| "img": base_img, | |
| "label": base_lab, | |
| "inst_map": base_inst_map, | |
| } | |
| ) | |
| return items | |
| def _save_mask(self, items, mask): | |
| mask = np.array(mask)[:,:,0] | |
| mask = mask.reshape((1,) + mask.shape).astype(np.float32) | |
| save_path = items["img_path"].replace(items["dataset"], "test_processed").replace("images", "predefined_masks/type_0").replace("jpg", "png") | |
| cv2.imwrite(save_path, mask[0]* 255) | |
| return mask[0].astype(np.uint8) | |
| def _edit_maps(self, items, mask, label, save=False): | |
| mask_path = items["img_path"].replace(items["dataset"], "test_processed").replace("images", "predefined_masks/type_0").replace("jpg", "png") | |
| mask = cv2.imread(mask_path, 0) / 255 | |
| target_pixels = mask == 1 | |
| target_inst_id = AppUtils.get_inst_id(self._input_id, label) | |
| items["inst_map"][target_pixels] = target_inst_id | |
| items["label"][target_pixels] = (target_inst_id % 120) | |
| im = Image.fromarray(items["inst_map"]).convert("I") | |
| if save: | |
| im.save(items["inst_map_path"].replace(items["dataset"], "test_processed")) | |
| cv2.imwrite(items["label_path"].replace(items["dataset"], "test_processed"), items["label"]) | |