Spaces:
Runtime error
Runtime error
| """ | |
| tree-segmentation | |
| Proof of concept showing effectiveness of a fine tuned instance segmentation model for detecting trees. | |
| """ | |
| import os | |
| import cv2 | |
| os.system("pip install 'git+https://github.com/facebookresearch/detectron2.git'") | |
| from transformers import DetrFeatureExtractor, DetrForSegmentation | |
| from PIL import Image | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| import torchvision | |
| import detectron2 | |
| # import some common detectron2 utilities | |
| import itertools | |
| import seaborn as sns | |
| from detectron2 import model_zoo | |
| from detectron2.engine import DefaultPredictor | |
| from detectron2.config import get_cfg | |
| from detectron2.utils.visualizer import Visualizer | |
| from detectron2.utils.visualizer import ColorMode | |
| from detectron2.data import MetadataCatalog, DatasetCatalog | |
| from detectron2.checkpoint import DetectionCheckpointer | |
| from detectron2.utils.visualizer import ColorMode | |
| from detectron2.structures import Instances | |
| # Model for trees | |
| tree_cfg = get_cfg() | |
| tree_cfg.merge_from_file("tree_model_weights/tree_cfg.yml") | |
| tree_cfg.MODEL.DEVICE='cpu' | |
| tree_cfg.MODEL.WEIGHTS = "tree_model_weights/treev1_best.pth" | |
| tree_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2 | |
| tree_predictor = DefaultPredictor(tree_cfg) | |
| # Model for buildings | |
| building_cfg = get_cfg() | |
| building_cfg.merge_from_file("building_model_weight/buildings_poc_cfg.yml") | |
| building_cfg.MODEL.DEVICE='cpu' | |
| building_cfg.MODEL.WEIGHTS = "building_model_weight/model_final.pth" | |
| building_cfg.MODEL.ROI_HEADS.NUM_CLASSES = 8 | |
| building_predictor = DefaultPredictor(building_cfg) | |
| # A function that runs the buildings model on an given image and confidence threshold | |
| def segment_building(im, confidence_threshold): | |
| building_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold | |
| outputs = building_predictor(im) | |
| building_instances = outputs["instances"].to("cpu") | |
| return building_instances | |
| # A function that runs the trees model on an given image and confidence threshold | |
| def segment_tree(im, confidence_threshold): | |
| tree_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold | |
| outputs = tree_predictor(im) | |
| tree_instances = outputs["instances"].to("cpu") | |
| return tree_instances | |
| # Function to map strings to color mode | |
| def map_color_mode(color_mode): | |
| if color_mode == "Black/white": | |
| return ColorMode.IMAGE_BW | |
| elif color_mode == "Random": | |
| return ColorMode.IMAGE | |
| elif color_mode == "Segmentation" or color_mode == None: | |
| return ColorMode.SEGMENTATION | |
| def visualize_image(im, mode, tree_threshold:float, building_threshold:float, color_mode): | |
| im = np.array(im) | |
| color_mode = map_color_mode(color_mode) | |
| if mode == "Trees": | |
| instances = segment_tree(im, tree_threshold) | |
| elif mode == "Buildings": | |
| instances = segment_building(im, building_threshold) | |
| elif mode == "Both" or mode == None: | |
| tree_instances = segment_tree(im, tree_threshold) | |
| building_instances = segment_building(im, building_threshold) | |
| instances = Instances.cat([tree_instances, building_instances]) | |
| metadata = MetadataCatalog.get("urban-trees-fdokv_train") | |
| print("metadata", type(metadata), metadata) | |
| print('metadata.get("thing_classes")', type(metadata.get("thing_classes")), metadata.get("thing_classes")) | |
| visualizer = Visualizer(im[:, :, ::-1], | |
| metadata=metadata, | |
| scale=0.5, | |
| instance_mode=color_mode) | |
| dataset_names = MetadataCatalog.list() | |
| print(dataset_names) | |
| metadata = MetadataCatalog.get("urban-small_train") | |
| category_names = metadata.get("thing_classes") | |
| print(category_names) | |
| # visualizer = Visualizer(im[:, :, ::-1], | |
| # metadata=metadata, | |
| # scale=0.5, | |
| # instance_mode=color_mode) | |
| # # in the visualizer, add category label names to detected instances | |
| # for instance in instances: | |
| # label = category_names[instance["category_id"]] | |
| # visualizer.draw_text(label, instance["bbox"][:2]) | |
| output_image = visualizer.draw_instance_predictions(instances) | |
| return Image.fromarray(output_image.get_image()[:, :, ::-1]) | |