Arulkumar03 commited on
Commit
a45382c
·
1 Parent(s): ea53030

Upload 2 files

Browse files
Files changed (2) hide show
  1. app2.py +54 -0
  2. requirements.txt +5 -0
app2.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import gradio as gr
4
+ from detectron2 import model_zoo
5
+ from detectron2.config import get_cfg
6
+ from detectron2.engine import DefaultPredictor
7
+ from detectron2.utils.visualizer import Visualizer
8
+ from detectron2.data import MetadataCatalog
9
+
10
+ def initialize_model():
11
+ for d in ["train", "test"]:
12
+ #DatasetCatalog.register("Animals_" + d, lambda d=d: get_wheat_dicts("Animal_Detection/" + d))
13
+ MetadataCatalog.get("Animals_" + d).set(thing_classes=["fox","sheep"])
14
+
15
+ wheat_metadata = MetadataCatalog.get("Animals_train")
16
+ cfg = get_cfg()
17
+ cfg.MODEL.DEVICE = "cpu"
18
+ cfg.DATALOADER.NUM_WORKERS = 0
19
+ cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_101_C4_3x.yaml")
20
+ cfg.SOLVER.IMS_PER_BATCH = 2
21
+ cfg.SOLVER.BASE_LR = 0.00025
22
+ cfg.SOLVER.STEPS = []
23
+ cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
24
+ cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
25
+ cfg.MODEL.WEIGHTS = "output/model_final.pth"
26
+ cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.95
27
+ predictor = DefaultPredictor(cfg)
28
+ return predictor
29
+
30
+ def process_image(predictor, img):
31
+ outputs = predictor(img)
32
+ wheat_metadata = MetadataCatalog.get("Animals_train")
33
+ v = Visualizer(img[:, :, ::-1],
34
+ metadata=wheat_metadata,
35
+ scale=1.5,
36
+ instance_mode="segmentation")
37
+ out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
38
+ processed_img = cv2.cvtColor(out.get_image()[:, :, ::-1], cv2.COLOR_BGR2RGB)
39
+ return processed_img
40
+
41
+ def main(img):
42
+ predictor = initialize_model()
43
+ processed_img = process_image(predictor, img)
44
+ return processed_img
45
+
46
+
47
+ iface = gr.Interface(
48
+ fn=main,
49
+ inputs="image",
50
+ outputs="image",
51
+ title="Fox & Sheep Computer Vision detector",
52
+ cache_examples=False,input_size=(5000, 5000), output_size=(8000, 8000)
53
+ )
54
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ pyyaml==5.1
2
+ torch==1.8.0
3
+ torchvision==0.9.0
4
+ git+https://github.com/facebookresearch/detectron2.git detectron2
5
+ opencv-contrib-python-headless