Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| import tree_commons as tc | |
| from PIL import Image | |
| import os | |
| import json | |
| import math | |
| import random | |
| import albumentations as A | |
| import numpy as np | |
| import cv2 | |
| import shutil | |
| import sahi | |
| from sahi.predict import get_sliced_prediction | |
| class ImageAugmenter: | |
| def __init__(self): | |
| self.noise = np.random.randint(0, 256, (tc.CROPPED_IMAGE_HEIGHT, tc.CROPPED_IMAGE_WIDTH, 3), dtype=np.uint8) | |
| def _noise_indivi_trees_randomly(self, img_arr, mask, p_indivi_survival=1 / 3): | |
| p_indivi_kill = 1 - p_indivi_survival | |
| instance_ids = np.unique(mask) | |
| indivi_kill_no = int(p_indivi_kill * len(instance_ids)) | |
| if indivi_kill_no == 0: | |
| return img_arr, mask | |
| indivi_tree_ids = instance_ids[instance_ids % 2 == 1] | |
| random.shuffle(indivi_tree_ids) | |
| not_survived_indivi_trees = indivi_tree_ids[:indivi_kill_no] | |
| not_survived_trees_mask_bool = np.isin(mask, not_survived_indivi_trees) | |
| not_survived_trees_mask = (not_survived_trees_mask_bool > 0).astype(np.uint8) | |
| survived_trees_mask = mask * (1 - not_survived_trees_mask) | |
| not_survived_trees_mask_3c = cv2.cvtColor(not_survived_trees_mask, cv2.COLOR_GRAY2RGB) | |
| blended = img_arr * (1 - not_survived_trees_mask_3c) + self.noise * not_survived_trees_mask_3c | |
| return blended, survived_trees_mask | |
| def _add_shadows(self, img_arr, mask): | |
| dis = [-25, -20, -10, -5, 5, 10, 20, 25] | |
| dx = np.random.choice(dis) | |
| dy = np.random.choice(dis) | |
| M = np.float32([[1, 0, dx], [0, 1, dy]]) # type: ignore | |
| shadow_mask = (mask > 0).astype(np.uint8) | |
| shadow_mask = cv2.warpAffine(shadow_mask, M, (mask.shape[1], mask.shape[0])) # type: ignore | |
| shadow_mask = cv2.GaussianBlur(shadow_mask, (13, 13), 70) | |
| mask_3c = cv2.cvtColor(shadow_mask, cv2.COLOR_GRAY2RGB) | |
| intensity = np.random.choice([0.20, 0.25, 0.30, 0.45,0.50]) | |
| shadowed_arr = img_arr * (1 - intensity * mask_3c) | |
| return np.clip(shadowed_arr, 0, 255).astype(np.uint8) | |
| def augment_image(self, image_arr, mask): | |
| h, w = tc.CROPPED_IMAGE_HEIGHT, tc.CROPPED_IMAGE_WIDTH | |
| transform = A.Compose([ | |
| A.RandomCrop(height=h, width=w, p=1.0, fill_mask=0), | |
| A.SquareSymmetry(p=0.1), | |
| A.RandomBrightnessContrast( | |
| brightness_limit=(-0.15, 0.15), | |
| contrast_limit=(-0.05, 0.05), | |
| p=0.10 | |
| ), | |
| A.MaskDropout(max_objects=(2, 10), fill=0.0, fill_mask=0.0, p=0.15), | |
| A.Affine(scale=(0.90, 1.1), keep_ratio=True, p = 0.15), | |
| A.CoarseDropout(num_holes_range=(1, 3), fill_mask=0, p=0.10), | |
| A.GaussianBlur(blur_limit=(3, 5), p=0.20), | |
| A.OneOf([ | |
| A.ToGray(p=1.0), | |
| A.ChannelDropout(p=0.2) | |
| ], p=0.20), | |
| ]) | |
| out = transform(image=image_arr, mask=mask) | |
| aug_img_arr, aug_mask = out["image"], out["mask"] | |
| if np.random.uniform(0, 1) <= 0.50: | |
| survival = np.random.choice([0.5/3, 1/3, 1.5/3]) | |
| aug_img_arr, aug_mask = self._noise_indivi_trees_randomly(aug_img_arr, aug_mask, survival) | |
| if np.random.uniform(0, 1) <= 0.50: | |
| aug_img_arr = self._add_shadows(aug_img_arr, aug_mask) | |
| return aug_img_arr, aug_mask | |
| def random_crop_image(self, image_arr, mask): | |
| h, w = tc.CROPPED_IMAGE_HEIGHT, tc.CROPPED_IMAGE_WIDTH | |
| transform = A.Compose([ | |
| A.RandomCrop(height=h, width=w, p=1.0, fill_mask=0), | |
| A.SquareSymmetry(p=0.25) | |
| ]) | |
| out = transform(image=image_arr, mask=mask) | |
| cropped_img, cropped_mask = out["image"], out["mask"] | |
| return cropped_img, cropped_mask | |
| class ImageDataCache: | |
| def __init__(self, no_samples_for_val = 2): | |
| self._filename_to_img_data = self._get_image_data_by_filename() | |
| self._file_type_to_filenames = self._get_filenames_by_file_type(self._filename_to_img_data.keys()) | |
| self._no_samples_for_val = no_samples_for_val | |
| def _get_filenames_by_file_type(self, filenames): | |
| file_type_to_filenames = {} | |
| for prefix in tc.FILE_TYPES_PREFIX: | |
| file_type_to_filenames[prefix] = [] | |
| for filename in filenames: | |
| for prefix in tc.FILE_TYPES_PREFIX: | |
| if filename.startswith(prefix): | |
| file_type_to_filenames[prefix].append(filename) | |
| break | |
| return file_type_to_filenames | |
| def sample_files(self, no, percents, train): | |
| sampled_files = [] | |
| for file_type, filenames in self._file_type_to_filenames.items(): | |
| sorted_filenames = sorted(filenames) | |
| if train: | |
| filesubset = sorted_filenames[0:] | |
| else: | |
| filesubset = sorted_filenames[-self._no_samples_for_val:] | |
| count = int(math.ceil(percents[file_type] * no)) | |
| sampled_files.extend(random.choices(filesubset, k=count)) | |
| return sampled_files | |
| def _get_image_data_by_filename(self): | |
| H, W = tc.IMAGE_HEIGHT, tc.IMAGE_WIDTH | |
| with open(tc.TRAIN_ANNOTATIONS_PATH, 'r') as file: | |
| data = json.load(file) | |
| filename_to_img_data = {} | |
| images_data = data[tc.IMAGES_KEY] | |
| for image_data in images_data: | |
| filename = image_data[tc.FILENAME_KEY] | |
| img_arr = tc.get_train_image_arr(filename) | |
| indivi_tree_instance_id = 1 | |
| grp_tree_instance_id = 2 | |
| mask = np.zeros((H, W), np.uint16) | |
| annotations = image_data[tc.ANNOTATIONS_KEY] | |
| for annotation in annotations: | |
| polygon = tc.get_polygon(annotation[tc.SEGMENTATION_KEY]) | |
| if annotation[tc.CLASS_KEY] == tc.CLASS_INDIVIDUAL_TREE: | |
| instance_id = indivi_tree_instance_id | |
| indivi_tree_instance_id += 2 | |
| else: | |
| instance_id = grp_tree_instance_id | |
| grp_tree_instance_id += 2 | |
| cv2.fillPoly(mask, [polygon], instance_id) # type: ignore | |
| filename_to_img_data[filename] = (img_arr, mask) | |
| return filename_to_img_data | |
| def get_image_data(self, filename): | |
| return self._filename_to_img_data[filename] | |
| def get_contours(mask): | |
| all_contours = [] | |
| instance_ids = np.unique(mask) | |
| instance_ids = instance_ids[1:] | |
| for instance_id in instance_ids: | |
| mask_bin = (mask == instance_id).astype(np.uint8) | |
| contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| if contours: | |
| all_contours.append(contours[0]) | |
| return all_contours | |
| def get_class_to_polygons_map(img_arr, mask): | |
| H = tc.CROPPED_IMAGE_HEIGHT | |
| W = tc.CROPPED_IMAGE_WIDTH | |
| indiv_tree_mask = mask * (mask % 2 == 1).astype(np.uint8) | |
| indiv_tree_contours = get_contours(indiv_tree_mask) | |
| grp_tree_mask = mask * (mask % 2 == 0).astype(np.uint8) | |
| grp_tree_contours = get_contours(grp_tree_mask) | |
| indiv_tree_class_id = tc.LABEL_TO_CLASS_ID[tc.CLASS_INDIVIDUAL_TREE] | |
| grp_tree_class_id = tc.LABEL_TO_CLASS_ID[tc.CLASS_GROUP_TREES] | |
| class_to_polygons_map = {indiv_tree_class_id :[], grp_tree_class_id:[]} | |
| for contours, class_id in zip([indiv_tree_contours, grp_tree_contours], [indiv_tree_class_id, grp_tree_class_id]): | |
| for contour in contours: | |
| polygon = contour.squeeze(1).astype(np.float32) | |
| polygon[:,0] /= W | |
| polygon[:,1] /= H | |
| class_to_polygons_map[class_id].append(polygon) | |
| return class_to_polygons_map | |
| class ImageDataset(): | |
| def __init__(self, image_data_cache, no, percents, train=True): | |
| self._no = no | |
| self._img_augmenter = ImageAugmenter() | |
| self._image_data_cache = image_data_cache | |
| self._percents = percents | |
| self._train = train | |
| self._filenames = self._image_data_cache.sample_files(self._no, self._percents, train) | |
| def __len__(self): | |
| return self._no | |
| def set_image_in_yolo_input_dir(self, indx): | |
| filename = self._filenames[indx] | |
| img_arr, mask = self._image_data_cache.get_image_data(filename) | |
| if self._train: | |
| aug_img_arr, aug_mask = self._img_augmenter.augment_image(img_arr, mask) | |
| image_output_dir = tc.get_yolo11_train_images_dir() | |
| labels_output_dir = tc.get_yolo11_train_labels_dir() | |
| else: | |
| aug_img_arr, aug_mask = self._img_augmenter.random_crop_image(img_arr, mask) | |
| image_output_dir = tc.get_yolo11_val_images_dir() | |
| labels_output_dir = tc.get_yolo11_val_labels_dir() | |
| class_to_polygons_map = get_class_to_polygons_map(aug_img_arr, aug_mask) | |
| image_id = indx | |
| image_name = tc.get_image_name_from_file_name(filename) | |
| image_output_name = f'{image_name}_{image_id}' | |
| image_output_filename = f'{image_output_dir}/{image_output_name}.png' | |
| os.makedirs(image_output_dir, exist_ok=True) | |
| img = Image.fromarray(aug_img_arr) | |
| img.save(image_output_filename) | |
| labels_output_filename = f'{labels_output_dir}/{image_output_name}.txt' | |
| txt_string = [] | |
| for class_id, polygons in class_to_polygons_map.items(): | |
| for polygon in polygons: | |
| str_polygons = ' '.join([f'{x} {y}' for x,y in polygon]) | |
| txt_string.append(f'{class_id} {str_polygons}') | |
| string = '\n'.join(txt_string) | |
| os.makedirs(labels_output_dir, exist_ok=True) | |
| with open(labels_output_filename, 'w') as f: | |
| f.write(string) | |
| def train(): | |
| if os.path.isdir(tc.YOLO11_INPUT_DIR): | |
| shutil.rmtree(tc.YOLO11_INPUT_DIR) | |
| train_percents = {tc.PREFIX_10CM_FILE_TYPE: 0.15, tc.PREFIX_20CM_FILE_TYPE: 0.15, | |
| tc.PREFIX_40CM_FILE_TYPE: 0.20, tc.PREFIX_60CM_FILE_TYPE: 0.25, | |
| tc.PREFIX_80CM_FILE_TYPE: 0.25} | |
| validation_percents = {tc.PREFIX_10CM_FILE_TYPE: 0.20, tc.PREFIX_20CM_FILE_TYPE: 0.20, | |
| tc.PREFIX_40CM_FILE_TYPE: 0.20, tc.PREFIX_60CM_FILE_TYPE: 0.20, | |
| tc.PREFIX_80CM_FILE_TYPE: 0.20} | |
| image_data_cache = ImageDataCache() | |
| train_dataset = ImageDataset(image_data_cache, 2500, train_percents, train=True) | |
| valid_dataset = ImageDataset(image_data_cache, 180, validation_percents, train=False) | |
| for indx in range(len(train_dataset)): | |
| train_dataset.set_image_in_yolo_input_dir(indx) | |
| for indx in range(len(valid_dataset)): | |
| valid_dataset.set_image_in_yolo_input_dir(indx) | |
| model = YOLO('yolo11x-seg.pt') | |
| res = model.train(data='data.yaml', epochs=100, batch=8, imgsz=tc.CROPPED_IMAGE_HEIGHT, save_period=10, cache=True, | |
| device=0, workers=8, name='medium_run', project=tc.YOLO11_OUTPUT_DIR, exist_ok=True, | |
| max_det=350, mask_ratio=2, classes=[0, 1], close_mosaic = 20, | |
| mosaic=0.20, hsv_s=0.10, hsv_h=0.10, translate=0.0) | |
| print(res) | |
| def predict(run_on_a_subset_of_eval_images): | |
| sahi_model = sahi.AutoDetectionModel.from_pretrained(model_type='ultralytics', | |
| model_path=tc.YOLO_BEST_WEIGHT, | |
| config_path='args.yaml', | |
| confidence_threshold=0.10, device='cpu', | |
| image_size = tc.CROPPED_IMAGE_HEIGHT) | |
| files = [] | |
| for filename in os.listdir(tc.EVALUATION_IMAGES_PATH): | |
| if filename.startswith('60cm') or not run_on_a_subset_of_eval_images: | |
| files.append(filename) | |
| results = [] | |
| for indx, filename in enumerate(files): | |
| print(f'remaining is {len(files) - indx}') | |
| filepath = os.path.join(tc.EVALUATION_IMAGES_PATH, filename) | |
| result = get_sliced_prediction( | |
| image=filepath, | |
| detection_model=sahi_model, | |
| slice_height=tc.CROPPED_IMAGE_HEIGHT, | |
| slice_width=tc.CROPPED_IMAGE_WIDTH, | |
| overlap_height_ratio=0.2, | |
| overlap_width_ratio=0.2, | |
| ) | |
| results.append(result) | |
| for filename, result in zip(files, results): | |
| annotations = [] | |
| for ann in result.to_coco_predictions(): | |
| annotation = {tc.CLASS_KEY:ann['category_name'], 'confidence_score': ann['score'], tc.SEGMENTATION_KEY:ann['segmentation'][0]} | |
| annotations.append(annotation) | |
| image_data = {tc.ANNOTATIONS_KEY : annotations} | |
| tc.show_img_with_annotations(tc.get_evaluation_image_arr(filename), image_data) | |
| # train() | |
| predict(True) | |