|
try: |
|
import detectron2 |
|
except: |
|
import os |
|
os.system('pip install git+https://github.com/facebookresearch/detectron2.git') |
|
|
|
from matplotlib.pyplot import axis |
|
import gradio as gr |
|
import requests |
|
import numpy as np |
|
from torch import nn |
|
import requests |
|
|
|
import torch |
|
import detectron2 |
|
from detectron2 import model_zoo |
|
from detectron2.engine import DefaultPredictor |
|
from detectron2.config import get_cfg |
|
from detectron2.utils.visualizer import Visualizer |
|
from detectron2.data import MetadataCatalog |
|
from detectron2.utils.visualizer import ColorMode |
|
|
|
model_path = 'model_final.pth' |
|
|
|
cfg = get_cfg() |
|
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) |
|
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.6 |
|
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 |
|
cfg.MODEL.WEIGHTS = model_path |
|
|
|
if not torch.cuda.is_available(): |
|
cfg.MODEL.DEVICE='cpu' |
|
|
|
predictor = DefaultPredictor(cfg) |
|
my_metadata = MetadataCatalog.get("car_dataset_val") |
|
my_metadata.thing_classes = ["damage"] |
|
|
|
def merge_segment(pred_segm): |
|
merge_dict = {} |
|
for i in range(len(pred_segm)): |
|
merge_dict[i] = [] |
|
for j in range(i+1,len(pred_segm)): |
|
if torch.sum(pred_segm[i]*pred_segm[j])>0: |
|
merge_dict[i].append(j) |
|
|
|
to_delete = [] |
|
for key in merge_dict: |
|
for element in merge_dict[key]: |
|
to_delete.append(element) |
|
|
|
for element in to_delete: |
|
merge_dict.pop(element,None) |
|
|
|
empty_delete = [] |
|
for key in merge_dict: |
|
if merge_dict[key] == []: |
|
empty_delete.append(key) |
|
|
|
for element in empty_delete: |
|
merge_dict.pop(element,None) |
|
|
|
for key in merge_dict: |
|
for element in merge_dict[key]: |
|
pred_segm[key]+=pred_segm[element] |
|
|
|
except_elem = list(set(to_delete)) |
|
|
|
new_indexes = list(range(len(pred_segm))) |
|
for elem in except_elem: |
|
new_indexes.remove(elem) |
|
|
|
return pred_segm[new_indexes] |
|
|
|
def inference(image): |
|
print(image.height) |
|
|
|
height = image.height |
|
|
|
|
|
img = np.array(image) |
|
outputs = predictor(img) |
|
out_dict = outputs["instances"].to("cpu").get_fields() |
|
new_inst = detectron2.structures.Instances((1024,1024)) |
|
new_inst.set('pred_masks',merge_segment(out_dict['pred_masks'])) |
|
v = Visualizer(img[:, :, ::-1], |
|
metadata=my_metadata, |
|
scale=0.5, |
|
instance_mode=ColorMode.SEGMENTATION |
|
) |
|
|
|
|
|
out = v.draw_instance_predictions(new_inst) |
|
|
|
return out.get_image()[:, :, ::-1] |
|
|
|
title = "Detectron2 Car damage Detection" |
|
description = "This demo introduces an interactive playground for our trained Detectron2 model." |
|
|
|
gr.Interface( |
|
inference, |
|
[gr.inputs.Image(type="pil", label="Input")], |
|
gr.outputs.Image(type="numpy", label="Output"), |
|
title=title, |
|
description=description, |
|
examples=[]).launch() |
