add app

main.py

@@ -0,0 +1,354 @@
+from detectron2.utils.logger import setup_logger
+
+import streamlit as st
+
+setup_logger()
+
+import collections
+
+import torch, torchvision
+
+from detectron2 import model_zoo
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.data import MetadataCatalog
+from detectron2.structures import masks
+
+from detectron2.evaluation import COCOEvaluator, inference_on_dataset
+from detectron2.data import build_detection_test_loader
+
+import detectron2.data.transforms as T
+from detectron2.data import DatasetMapper   # the default mapper
+
+import os
+import cv2
+from matplotlib import pyplot as plt
+
+from IPython.display import display
+from PIL import Image
+
+
+# Some basic setup:
+# Setup detectron2 logger
+import detectron2
+from detectron2.utils.logger import setup_logger
+setup_logger()
+
+# import some common libraries
+import numpy as np
+import os, json, cv2, random
+import copy
+# from google.colab.patches import cv2_imshow
+
+# import some common detectron2 utilities
+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.datasets import register_coco_instances
+from detectron2.engine import DefaultTrainer
+
+from detectron2.evaluation import COCOEvaluator, inference_on_dataset
+from detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_test_loader, build_detection_train_loader
+
+from detectron2.data import transforms as T
+from detectron2.data import detection_utils as utils
+
+
+def predictor_define(cfg, CONFIDENCE=0.7):
+    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  # path to the model we just trained
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = CONFIDENCE  # set a custom testing threshold 0.7
+    predictor = DefaultPredictor(cfg)
+
+
+def inference(image, THR=0.5, path_to_model="/content/output/model_final.pth"):
+    cfg = get_cfg()
+    cfg.MODEL.DEVICE = "cuda"  # 'cpu' or 'cuda'
+
+    out_dct = dict()
+
+    cat_lst = ["0.2", "0.3", "0.5", "0.8"]
+
+    cfg.merge_from_file(
+        model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
+    )
+    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = THR
+    cfg.MODEL.WEIGHTS = path_to_model
+    MetadataCatalog.get("dataset").thing_classes = cat_lst
+
+    predictor = DefaultPredictor(cfg)
+    outputs = predictor(image)
+
+    box_lst = [i.tolist() for i in outputs["instances"].pred_boxes]
+    cl_lst = outputs["instances"].pred_classes.tolist()
+    score_lst = outputs["instances"].scores.tolist()
+
+    out_dct["box_lst"] = box_lst
+    out_dct["cl_lst"] = [cat_lst[i] for i in cl_lst]
+    out_dct["scores"] = score_lst
+
+    return out_dct
+
+
+def inference_2(image, path_to_model, dataset_name, YAML_FILE, cat_lst=["0.2", "0.3", "0.5", "0.8"], thr=0.25,
+                aa_dct='None', device='cuda', loader='default', aug=[            T.Resize((512, 512))        ]):
+    cfg = get_cfg()
+    cfg.MODEL.DEVICE = device
+
+    out_dct = dict()
+
+    cfg.merge_from_file(
+        model_zoo.get_config_file(YAML_FILE)
+    )
+    cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(cat_lst)
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = thr
+    cfg.INPUT.MIN_SIZE_TEST = 900
+    cfg.INPUT.MAX_SIZE_TEST = 1000
+    cfg.MODEL.WEIGHTS = path_to_model
+
+    if aa_dct != 'None':
+        cfg.MODEL.ANCHOR_GENERATOR.SIZES = aa_dct['anchor_sizes']
+        cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS = aa_dct['aspect_ratio']
+
+    # MetadataCatalog.get(dataset_name).thing_classes = cat_lst
+    predictor = DefaultPredictor(cfg)
+
+    # if loader == 'default':
+    #   val_loader = build_detection_test_loader(cfg, dataset_name)
+    # else:
+    #   val_loader = build_detection_test_loader(cfg, dataset_name, mapper=DatasetMapper(cfg, is_train=False, augmentations = aug))
+
+    # return_inference = inference_on_dataset(predictor.model, val_loader, evaluator)
+
+    outputs = predictor(image)
+
+    box_lst = [i.tolist() for i in outputs["instances"].pred_boxes]
+    cl_lst = outputs["instances"].pred_classes.tolist()
+    score_lst = outputs["instances"].scores.tolist()
+
+    out_dct["box_lst"] = box_lst
+    out_dct["cl_lst"] = [cat_lst[i] for i in cl_lst]
+    out_dct["scores"] = score_lst
+
+    return out_dct
+
+
+def imagename(str):
+    return str.split('/')[-1]
+
+
+def read_image(im_path, YAML_FILE, path_to_model, THR=0.5, dim=(500, 500),
+                                cat_lst=['pomeranc', 'poteklina']):
+
+    out_dct = {}
+
+    im_name = imagename(im_path)
+    # img = cv2.imread(os.path.join(root_path, d["file_name"]))
+    img = cv2.imread(im_path)
+    out_dct[im_name] = inference_2(image=img, path_to_model=path_to_model,
+                                   dataset_name=None, YAML_FILE=YAML_FILE,
+                                   cat_lst=cat_lst, thr=THR)
+    print(im_path)
+    # print(os.path.join(im_path, d["file_name"]))
+    labels_list = []
+    boxes = []
+    name_dict = {}
+
+    for i in range(len(out_dct[im_name]['cl_lst'])):
+        box = out_dct[im_name]['box_lst'][i]
+        box = [int(i) for i in box]
+        label = out_dct[im_name]['cl_lst'][i]
+        scores = out_dct[im_name]['scores'][i]
+        # print(label, scores)
+        # print(box)
+        labels_list.append(label)
+        boxes.append(box[0])
+
+        cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 1)
+        cv2.putText(img, label, (box[0], box[3]), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 4)
+    # cv2.putText(img, str(int(scores*100)), (box[0], box[3]),cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
+    # print(d["file_name"])
+    # print(labels_list)
+
+    for num, name in zip(boxes, labels_list):
+        name_dict[num] = name
+
+    # print(name_dict)
+
+    od = collections.OrderedDict(sorted(name_dict.items()))
+    digits_out_sorted = []
+    for k, v in od.items():
+        digits_out_sorted.append(v)
+
+    # print(od)
+
+    print(digits_out_sorted)
+
+    # resized = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
+    im_output_path = ROOT_FOLDER + 'output_images/' + im_path.split('/')[-1]
+    print(im_output_path)
+    cv2.imwrite(im_output_path, img)
+
+def process_one_image(img, im_name, YAML_FILE, path_to_model, THR=0.5, dim=(500, 500),
+                                cat_lst=['1', '2', '3', '4', '5', '6', '7', '8', '9', '0']):
+
+    out_dct = {}
+
+    # im_name = imagename(im_path)
+    # img = cv2.imread(os.path.join(root_path, d["file_name"]))
+    # img = cv2.imread(im_path)
+    out_dct[im_name] = inference_2(image=img, path_to_model=path_to_model,
+                                   dataset_name=None, YAML_FILE=YAML_FILE,
+                                   cat_lst=cat_lst, thr=THR)
+    # print(im_path)
+    # print(os.path.join(im_path, d["file_name"]))
+    labels_list = []
+    boxes = []
+    name_dict = {}
+
+    for i in range(len(out_dct[im_name]['cl_lst'])):
+        box = out_dct[im_name]['box_lst'][i]
+        box = [int(i) for i in box]
+        label = out_dct[im_name]['cl_lst'][i]
+        scores = out_dct[im_name]['scores'][i]
+        # print(label, scores)
+        # print(box)
+        labels_list.append(label)
+        boxes.append(box[0])
+
+        cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 1)
+        cv2.putText(img, label, (box[0], box[3]), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 4)
+    # cv2.putText(img, str(int(scores*100)), (box[0], box[3]),cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
+    # print(d["file_name"])
+    # print(labels_list)
+
+    for num, name in zip(boxes, labels_list):
+        name_dict[num] = name
+
+    # print(name_dict)
+
+    od = collections.OrderedDict(sorted(name_dict.items()))
+    digits_out_sorted = []
+    for k, v in od.items():
+        digits_out_sorted.append(v)
+
+    # print(od)
+
+    print(digits_out_sorted)
+
+    # resized = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
+    # im_output_path = ROOT_FOLDER + 'output_images/' + im_path.split('/')[-1]
+    # print(im_output_path)
+    # cv2.imwrite(im_output_path, img)
+
+    return img, digits_out_sorted
+
+ROOT_FOLDER = '/home/roman/PycharmProjects/streamlit/digits/'
+
+train_dct1 = {
+    '10cl':
+        {
+            'cat_lst': ['1', '2', '3', '4', '5', '6', '7', '8', '9', '0'],
+            'train': {
+                '13': {
+                    'json': ROOT_FOLDER + '/json/digits-13.json',
+                    'data': ROOT_FOLDER + '/data/train/',
+
+                },
+                '14': {
+                    'json': ROOT_FOLDER + '/json/digits-14.json',
+                    'data': ROOT_FOLDER + '/data/train/',
+
+                },
+                '17': {
+                    'json': ROOT_FOLDER + '/annotations/digits_04-17.json',
+                    'data': ROOT_FOLDER + '/images/train/',
+
+                },
+            },
+            'val': {
+                '16': {
+                    'json': ROOT_FOLDER + '/annotations/digits_02-16.json',
+                    'data': ROOT_FOLDER + '/images/val/',
+
+                },
+            }
+        }
+}
+
+if __name__ == '__main__':
+
+    st.title('DIGITS')
+
+    model_name = 'model_final'  # '30im_10cl_2000it_faster_rcnn_R_50_FPN_3x_None'
+    model_name = 'model_0003999'
+    # model_name = 'model_final_8000'
+    # model_name = 'model_final_4000_nms_01'
+    model_path = ROOT_FOLDER + 'output/' + model_name + '.pth'
+
+    YAML_FILE = 'Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml'
+    # YAML_FILE = 'COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml'
+
+    DATASET_NAME = model_name
+
+    val_dataset_name = DATASET_NAME + "_val_" + str(random.randint(1, 1000))
+
+    num_classes = '10cl'  # 'poteklina' # '1cl'
+    dataset = '16'  #
+
+    # MetadataCatalog.get(val_dataset_name).thing_classes = train_dct1['10cl']['cat_lst'] #["person", "dog"]
+    register_coco_instances(val_dataset_name, {}, train_dct1[num_classes]['val'][dataset]['json'],
+                            train_dct1[num_classes]['val'][dataset]['data'])
+    dataset_dicts = DatasetCatalog.get(val_dataset_name)
+
+    filenames = []
+    uploaded_files = st.file_uploader("Choose a images", accept_multiple_files=True, type=["png", "jpg", "jpeg"])
+
+    num_ok = 0
+    num_nok = 0
+    for uploaded_file in uploaded_files:
+
+        image = Image.open(uploaded_file).convert("RGB")
+        # Convert PIL image to array
+        image = np.array(image)
+
+        img, predicted_digits = process_one_image(image, uploaded_file.name, YAML_FILE=YAML_FILE, path_to_model=model_path, THR = 0.7)
+
+        print(uploaded_file.name)
+        print(predicted_digits)
+        st.image(image)
+
+        wrong_indexes = []
+        result_in_string = ''
+
+        for el in predicted_digits:
+            result_in_string += el
+
+        if len(predicted_digits) >= 6:
+            for num, el in enumerate(uploaded_file.name[:6]):
+
+                if el != predicted_digits[num]:
+                    wrong_indexes.append(num)
+                    print(el, num)
+        else:
+            wrong_indexes = ['0']
+
+        status = 'NOK'
+        if len(wrong_indexes) == 0:
+            status = 'OK'
+
+        if status == 'OK':
+            num_ok +=1
+        else:
+            num_nok +=1
+
+        file_details = {"filename": uploaded_file.name, "detection_result": result_in_string,
+         "status": status}
+
+        st.write(file_details)
+
+        st.write('Точність визначення = ', num_ok /(num_ok + num_nok) * 100, ' %')
+        st.write('NOK = ', num_nok, 'OK = ', num_ok)