Martin Tomov commited on
Commit
b547fbf
Β·
verified Β·
1 Parent(s): 123cd28

demo import

Browse files
Files changed (10) hide show
  1. README.md +23 -4
  2. app.py +33 -0
  3. demo.jpg +0 -0
  4. demo1.jpg +0 -0
  5. detectron_utils.py +130 -0
  6. gitattributes +35 -0
  7. gitattributes.txt +35 -0
  8. requirements.txt +11 -0
  9. sam_utils.py +224 -0
  10. yolo_utils.py +18 -0
README.md CHANGED
@@ -1,13 +1,32 @@
1
  ---
2
- title: InsectModelZoo
3
- emoji: ⚑
4
- colorFrom: red
5
  colorTo: green
6
  sdk: gradio
7
  sdk_version: 4.36.1
8
  app_file: app.py
9
  pinned: false
10
  license: apache-2.0
 
11
  ---
12
 
13
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ title: Insect Model Zoo
3
+ emoji: πŸžπŸ”¬
4
+ colorFrom: blue
5
  colorTo: green
6
  sdk: gradio
7
  sdk_version: 4.36.1
8
  app_file: app.py
9
  pinned: false
10
  license: apache-2.0
11
+ short_description: Comparing InsectSAM, Yolov8, Detectron2
12
  ---
13
 
14
+ ## https://insectsam.live
15
+
16
+ ## What is InsectSAM?
17
+
18
+ InsectSAM is an open-source machine learning model tailored for the DIOPSIS camera systems and ARISE algorithms, dedicated to Insect Biodiversity Detection and Monitoring in the Netherlands.
19
+
20
+ ![whatisinsectsam](https://insectsam.live/assets/images/undraw_docusaurus_mountain-e42b6f2eba6f6cca2d69178e66414779.png)
21
+
22
+ ## How does it work?
23
+
24
+ Based on Meta AI segment-anything, InsectSAM is fine-tuned to accurately segment insects from complex backgrounds. It boosts the precision and efficiency of biodiversity monitoring algorithms, especially in scenarios with diverse backgrounds that attract insects.
25
+
26
+ ![howdoesitwork](https://insectsam.live/assets/images/undraw_docusaurus_react-616141633fe960e3f853b86fef751af9.png)
27
+
28
+ ## Technologies Used
29
+
30
+ Python, PyTorch, Hugging Face Transformers, OpenCV, and more. InsectSAM is designed to be easily integrated into existing DIOPSIS and ARISE algorithms, providing a seamless experience for researchers and developers.
31
+
32
+ ![tech](https://insectsam.live/assets/images/undraw_docusaurus_stack-957770c8948b116dc310df4dbcf9bb34.png)
app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ try:
2
+ import detectron2
3
+ except:
4
+ import os
5
+ os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
6
+
7
+ import gradio as gr
8
+ import json
9
+ import numpy as np
10
+ from sam_utils import grounded_segmentation, create_yellow_background_with_insects
11
+ from yolo_utils import yolo_processimage
12
+ from detectron_utils import detectron_process_image
13
+ def process_image(image, include_json):
14
+ detectron_result=detectron_process_image(image)
15
+ yolo_result = yolo_processimage(image)
16
+ insectsam_result = create_yellow_background_with_insects(image)
17
+
18
+ return insectsam_result, yolo_result, detectron_result
19
+
20
+ examples = [
21
+ ["demo.jpg"],
22
+ ["demo1.jpg"],
23
+ ]
24
+
25
+ gr.Interface(
26
+ fn=process_image,
27
+ inputs=[gr.Image(type="pil")],
28
+ outputs=[gr.Image(label='InsectSAM', type="numpy"),
29
+ gr.Image(label='Yolov8', type="numpy"),
30
+ gr.Image(label='Detectron', type="numpy")],
31
+ title="Insect Model Zoo πŸžπŸ”¬",
32
+ examples=examples
33
+ ).launch()
demo.jpg ADDED
demo1.jpg ADDED
detectron_utils.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import cv2
4
+ from huggingface_hub import hf_hub_download
5
+
6
+ REPO_ID = "idml/Detectron2-FasterRCNN_InsectDetect"
7
+ FILENAME = "model.pth"
8
+ FILENAME_CONFIG = "config.yml"
9
+
10
+
11
+ # Ensure you have the model file
12
+
13
+ import cv2
14
+ from detectron2.config import get_cfg
15
+ from detectron2.engine import DefaultPredictor
16
+ from detectron2.data import MetadataCatalog
17
+ from detectron2.utils.visualizer import Visualizer, ColorMode
18
+ import matplotlib.pyplot as plt
19
+
20
+
21
+ viz_classes = {'thing_classes': ['Acrididae',
22
+ 'Agapeta',
23
+ 'Agapeta hamana',
24
+ 'Animalia',
25
+ 'Anisopodidae',
26
+ 'Aphididae',
27
+ 'Apidae',
28
+ 'Arachnida',
29
+ 'Araneae',
30
+ 'Arctiidae',
31
+ 'Auchenorrhyncha indet.',
32
+ 'Baetidae',
33
+ 'Cabera',
34
+ 'Caenidae',
35
+ 'Carabidae',
36
+ 'Cecidomyiidae',
37
+ 'Ceratopogonidae',
38
+ 'Cercopidae',
39
+ 'Chironomidae',
40
+ 'Chrysomelidae',
41
+ 'Chrysopidae',
42
+ 'Chrysoteuchia culmella',
43
+ 'Cicadellidae',
44
+ 'Coccinellidae',
45
+ 'Coleophoridae',
46
+ 'Coleoptera',
47
+ 'Collembola',
48
+ 'Corixidae',
49
+ 'Crambidae',
50
+ 'Culicidae',
51
+ 'Curculionidae',
52
+ 'Dermaptera',
53
+ 'Diptera',
54
+ 'Eilema',
55
+ 'Empididae',
56
+ 'Ephemeroptera',
57
+ 'Erebidae',
58
+ 'Fanniidae',
59
+ 'Formicidae',
60
+ 'Gastropoda',
61
+ 'Gelechiidae',
62
+ 'Geometridae',
63
+ 'Hemiptera',
64
+ 'Hydroptilidae',
65
+ 'Hymenoptera',
66
+ 'Ichneumonidae',
67
+ 'Idaea',
68
+ 'Insecta',
69
+ 'Lepidoptera',
70
+ 'Leptoceridae',
71
+ 'Limoniidae',
72
+ 'Lomaspilis marginata',
73
+ 'Miridae',
74
+ 'Mycetophilidae',
75
+ 'Nepticulidae',
76
+ 'Neuroptera',
77
+ 'Noctuidae',
78
+ 'Notodontidae',
79
+ 'Object',
80
+ 'Opiliones',
81
+ 'Orthoptera',
82
+ 'Panorpa germanica',
83
+ 'Panorpa vulgaris',
84
+ 'Parasitica indet.',
85
+ 'Plutellidae',
86
+ 'Psocodea',
87
+ 'Psychodidae',
88
+ 'Pterophoridae',
89
+ 'Pyralidae',
90
+ 'Pyrausta',
91
+ 'Sepsidae',
92
+ 'Spilosoma',
93
+ 'Staphylinidae',
94
+ 'Stratiomyidae',
95
+ 'Syrphidae',
96
+ 'Tettigoniidae',
97
+ 'Tipulidae',
98
+ 'Tomoceridae',
99
+ 'Tortricidae',
100
+ 'Trichoptera',
101
+ 'Triodia sylvina',
102
+ 'Yponomeuta',
103
+ 'Yponomeutidae']}
104
+
105
+
106
+
107
+ def detectron_process_image(image):
108
+ cfg = get_cfg()
109
+
110
+
111
+ cfg.merge_from_file(hf_hub_download(repo_id=REPO_ID, filename=FILENAME_CONFIG))
112
+ cfg.MODEL.WEIGHTS = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
113
+ cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2
114
+ cfg.MODEL.DEVICE='cpu'
115
+ predictor = DefaultPredictor(cfg)
116
+
117
+ numpy_image = np.array(image)
118
+
119
+
120
+ im = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
121
+
122
+ v = Visualizer(im[:, :, ::-1],
123
+ viz_classes,
124
+ scale=0.5)
125
+ outputs = predictor(im)
126
+ out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
127
+ results = out.get_image()[:, :, ::-1]
128
+ rgb_image = cv2.cvtColor(results, cv2.COLOR_BGR2RGB)
129
+
130
+ return rgb_image
gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
gitattributes.txt ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
requirements.txt ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gradio
2
+ torchvision
3
+ transformers
4
+ opencv-python
5
+ Pillow
6
+ numpy
7
+ requests
8
+ matplotlib
9
+ ultralytics
10
+ onnxruntime
11
+ efficientnet
sam_utils.py ADDED
@@ -0,0 +1,224 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ import random
4
+ from dataclasses import dataclass
5
+ from typing import Any, List, Dict, Optional, Union, Tuple
6
+ import cv2
7
+ import torch
8
+ import requests
9
+ import numpy as np
10
+ from PIL import Image
11
+ import matplotlib.pyplot as plt
12
+ from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
13
+ import gradio as gr
14
+ import json
15
+
16
+
17
+ @dataclass
18
+ class BoundingBox:
19
+ xmin: int
20
+ ymin: int
21
+ xmax: int
22
+ ymax: int
23
+
24
+ @property
25
+ def xyxy(self) -> List[float]:
26
+ return [self.xmin, self.ymin, self.xmax, self.ymax]
27
+ @dataclass
28
+ class DetectionResult:
29
+ score: float
30
+ label: str
31
+ box: BoundingBox
32
+ mask: Optional[np.ndarray] = None
33
+
34
+ @classmethod
35
+ def from_dict(cls, detection_dict: Dict) -> 'DetectionResult':
36
+ return cls(
37
+ score=detection_dict['score'],
38
+ label=detection_dict['label'],
39
+ box=BoundingBox(
40
+ xmin=detection_dict['box']['xmin'],
41
+ ymin=detection_dict['box']['ymin'],
42
+ xmax=detection_dict['box']['xmax'],
43
+ ymax=detection_dict['box']['ymax']
44
+ )
45
+ )
46
+
47
+ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:
48
+ image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
49
+ image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
50
+
51
+ for detection in detection_results:
52
+ label = detection.label
53
+ score = detection.score
54
+ box = detection.box
55
+ mask = detection.mask
56
+
57
+ if include_bboxes:
58
+ color = np.random.randint(0, 256, size=3).tolist()
59
+ cv2.rectangle(image_cv2, (box.xmin, box.ymin),
60
+ (box.xmax, box.ymax), color, 2)
61
+ cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
62
+ cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
63
+
64
+ return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
65
+
66
+
67
+ def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:
68
+ annotated_image = annotate(image, detections, include_bboxes)
69
+ return annotated_image
70
+
71
+
72
+ def load_image(image: Union[str, Image.Image]) -> Image.Image:
73
+ if isinstance(image, str) and image.startswith("http"):
74
+ image = Image.open(requests.get(image, stream=True).raw).convert("RGB")
75
+ elif isinstance(image, str):
76
+ image = Image.open(image).convert("RGB")
77
+ else:
78
+ image = image.convert("RGB")
79
+ return image
80
+
81
+
82
+ def get_boxes(detection_results: List[DetectionResult]) -> List[List[List[float]]]:
83
+ boxes = []
84
+ for result in detection_results:
85
+ xyxy = result.box.xyxy
86
+ boxes.append(xyxy)
87
+ return [boxes]
88
+
89
+
90
+ def mask_to_polygon(mask: np.ndarray) -> np.ndarray:
91
+ contours, _ = cv2.findContours(
92
+ mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
93
+ if len(contours) == 0:
94
+ return np.array([])
95
+ largest_contour = max(contours, key=cv2.contourArea)
96
+ return largest_contour
97
+
98
+
99
+ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]:
100
+ masks = masks.cpu().float().permute(0, 2, 3, 1).mean(
101
+ axis=-1).numpy().astype(np.uint8)
102
+ masks = (masks > 0).astype(np.uint8)
103
+ if polygon_refinement:
104
+ for idx, mask in enumerate(masks):
105
+ shape = mask.shape
106
+ polygon = mask_to_polygon(mask)
107
+ masks[idx] = cv2.fillPoly(
108
+ np.zeros(shape, dtype=np.uint8), [polygon], 1)
109
+ return list(masks)
110
+
111
+
112
+ def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[Dict[str, Any]]:
113
+ detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
114
+ object_detector = pipeline(
115
+ model=detector_id, task="zero-shot-object-detection", device="cpu")
116
+ labels = [label if label.endswith(".") else label+"." for label in labels]
117
+ results = object_detector(
118
+ image, candidate_labels=labels, threshold=threshold)
119
+ return [DetectionResult.from_dict(result) for result in results]
120
+
121
+
122
+ def segment(image: Image.Image, detection_results: List[DetectionResult], polygon_refinement: bool = False, segmenter_id: Optional[str] = None) -> List[DetectionResult]:
123
+ segmenter_id = segmenter_id if segmenter_id else "martintmv/InsectSAM"
124
+ segmentator = AutoModelForMaskGeneration.from_pretrained(
125
+ segmenter_id).to("cpu")
126
+ processor = AutoProcessor.from_pretrained(segmenter_id)
127
+ boxes = get_boxes(detection_results)
128
+ inputs = processor(images=image, input_boxes=boxes,
129
+ return_tensors="pt").to("cpu")
130
+ outputs = segmentator(**inputs)
131
+ masks = processor.post_process_masks(
132
+ masks=outputs.pred_masks, original_sizes=inputs.original_sizes, reshaped_input_sizes=inputs.reshaped_input_sizes)[0]
133
+ masks = refine_masks(masks, polygon_refinement)
134
+ for detection_result, mask in zip(detection_results, masks):
135
+ detection_result.mask = mask
136
+ return detection_results
137
+
138
+
139
+ def grounded_segmentation(image: Union[Image.Image, str], labels: List[str], threshold: float = 0.3, polygon_refinement: bool = False, detector_id: Optional[str] = None, segmenter_id: Optional[str] = None) -> Tuple[np.ndarray, List[DetectionResult]]:
140
+ image = load_image(image)
141
+ detections = detect(image, labels, threshold, detector_id)
142
+ detections = segment(image, detections, polygon_refinement, segmenter_id)
143
+ return np.array(image), detections
144
+
145
+
146
+ def mask_to_min_max(mask: np.ndarray) -> Tuple[int, int, int, int]:
147
+ y, x = np.where(mask)
148
+ return x.min(), y.min(), x.max(), y.max()
149
+
150
+
151
+ def extract_and_paste_insect(original_image: np.ndarray, detection: DetectionResult, background: np.ndarray) -> None:
152
+ mask = detection.mask
153
+ xmin, ymin, xmax, ymax = mask_to_min_max(mask)
154
+ insect_crop = original_image[ymin:ymax, xmin:xmax]
155
+ mask_crop = mask[ymin:ymax, xmin:xmax]
156
+
157
+ insect = cv2.bitwise_and(insect_crop, insect_crop, mask=mask_crop)
158
+
159
+ x_offset, y_offset = xmin, ymin
160
+ x_end, y_end = x_offset + insect.shape[1], y_offset + insect.shape[0]
161
+
162
+ insect_area = background[y_offset:y_end, x_offset:x_end]
163
+ insect_area[mask_crop == 1] = insect[mask_crop == 1]
164
+
165
+
166
+ def create_yellow_background_with_insects(image: np.ndarray) -> np.ndarray:
167
+ labels = ["insect"]
168
+
169
+ original_image, detections = grounded_segmentation(
170
+ image, labels, threshold=0.3, polygon_refinement=True)
171
+
172
+ yellow_background = np.full(
173
+ (original_image.shape[0], original_image.shape[1], 3), (0, 255, 255), dtype=np.uint8) # BGR for yellow
174
+ for detection in detections:
175
+ if detection.mask is not None:
176
+ extract_and_paste_insect(
177
+ original_image, detection, yellow_background)
178
+ # Convert back to RGB to match Gradio's expected input format
179
+ yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
180
+ return yellow_background
181
+
182
+
183
+ def run_length_encoding(mask):
184
+ pixels = mask.flatten()
185
+ rle = []
186
+ last_val = 0
187
+ count = 0
188
+ for pixel in pixels:
189
+ if pixel == last_val:
190
+ count += 1
191
+ else:
192
+ if count > 0:
193
+ rle.append(count)
194
+ count = 1
195
+ last_val = pixel
196
+ if count > 0:
197
+ rle.append(count)
198
+ return rle
199
+
200
+
201
+ def detections_to_json(detections):
202
+ detections_list = []
203
+ for detection in detections:
204
+ detection_dict = {
205
+ "score": detection.score,
206
+ "label": detection.label,
207
+ "box": {
208
+ "xmin": detection.box.xmin,
209
+ "ymin": detection.box.ymin,
210
+ "xmax": detection.box.xmax
211
+ },
212
+ "mask": run_length_encoding(detection.mask) if detection.mask is not None else None
213
+ }
214
+ detections_list.append(detection_dict)
215
+ return detections_list
216
+
217
+
218
+ def crop_bounding_boxes_with_yellow_background(image: np.ndarray, yellow_background: np.ndarray, detections: List[DetectionResult]) -> List[np.ndarray]:
219
+ crops = []
220
+ for detection in detections:
221
+ xmin, ymin, xmax, ymax = detection.box.xyxy
222
+ crop = yellow_background[ymin:ymax, xmin:xmax]
223
+ crops.append(crop)
224
+ return crops
yolo_utils.py ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ultralytics import YOLO
2
+ import torch
3
+ import numpy as np
4
+ import cv2
5
+ from huggingface_hub import hf_hub_download
6
+
7
+ REPO_ID = "idml/Yolov8_InsectDetect"
8
+ FILENAME = "insectYolo.pt"
9
+
10
+
11
+ # Ensure you have the model file
12
+ model = YOLO(hf_hub_download(repo_id=REPO_ID, filename=FILENAME))
13
+ def yolo_processimage(image):
14
+ results = model(source=image,
15
+ conf=0.2, device='cpu')
16
+ rgb_image = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
17
+ return rgb_image
18
+