Spaces:
Paused
Paused
File size: 2,463 Bytes
3f7dd83 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import logging
from control_models.base import ControlModel
from typing import List, Dict
logger = logging.getLogger(__name__)
class PolygonLabelsModel(ControlModel):
"""
Class representing a PolygonLabels control tag for YOLO model.
"""
type = "PolygonLabels"
model_path = "yolov8n-seg.pt"
@classmethod
def is_control_matched(cls, control) -> bool:
# check object tag type
if control.objects[0].tag != "Image":
return False
return control.tag == cls.type
def predict_regions(self, path) -> List[Dict]:
results = self.model.predict(path)
return self.create_polygons(results, path)
def create_polygons(self, results, path):
logger.debug(f"create_polygons: {self.from_name}")
data = results[0].masks # take masks from the first frame
model_names = self.model.names
regions = []
for i in range(len(data)):
score = float(results[0].boxes.conf[i]) # tensor => float
points = (
data.xyn[i] * 100
) # get the polygon points for the current instance
model_label = model_names[int(results[0].boxes.cls[i])]
logger.debug(
"----------------------\n"
f"task id > {path}\n"
f"type: {self.control}\n"
f"polygon points > {points}\n"
f"model label > {model_label}\n"
f"score > {score}\n"
)
# bbox score is too low
if score < self.model_score_threshold:
continue
# there is no mapping between model label and LS label
if model_label not in self.label_map:
continue
output_label = self.label_map[model_label]
# add new region with polygon
region = {
"from_name": self.from_name,
"to_name": self.to_name,
"type": "polygonlabels",
"value": {
"polygonlabels": [output_label],
"points": points.tolist(), # Converting the tensor to a list for JSON serialization
"closed": True,
},
"score": score,
}
regions.append(region)
return regions
# pre-load and cache default model at startup
PolygonLabelsModel.get_cached_model(PolygonLabelsModel.model_path)
|