detectron2_for_windows/instance_segmentation.py at main · naokiwakata/detectron2_for_windows · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
from dataclasses import field
import cv2
from fileEnum import File
from processImage import processImage
from predictor import Predictor
import numpy as np
from util import shaveOff
import glob

def instanceSegmentation():
    imagePath = "images\9_23 (3).JPG"
    img = cv2.imread(imagePath)  # <class 'numpy.ndarray'>

    predictor = Predictor()

    outputs = predictor.predict(img=img)

    # jpgs = glob.glob('testDataLeaf\\*.jpg')
    # for imagePath in jpgs:
    #     img = cv2.imread(imagePath)  # <class 'numpy.ndarray'>
    #     outputs = predictor.predict(img=img)
    #     fields = outputs['instances'].get_fields()
    #     pred_boxes = fields['pred_boxes']
    #     print(len(pred_boxes))
    #     predictor.showPredictImage(img=img, outputs=outputs)


    shaveOff(outputs=outputs,img=img) #葉っぱのみを切り抜く

    fields = outputs['instances'].get_fields()
    pred_boxes = fields['pred_boxes']
    scores = fields['scores'].to('cpu').detach().numpy()
    pred_classes = fields['pred_classes'].to('cpu').detach().numpy()
    pred_masks = fields['pred_masks'].to('cpu').detach().numpy()

    image_size = outputs['instances'].image_size
    height = image_size[0]
    width = image_size[1]

    predictor.showPredictImage(img=img, outputs=outputs)

    file = File.Image
    #file = File.Video

    if file == File.Image:
        #processImage(outputs=outputs, img=img)
        print(file)
    elif file == File.Video:
        print(file)