label-studio-ml-backend
label-studio-ml-backend copied to clipboard
AttributeError: 'FasterRCNN' object has no attribute 'CLASSES'
hello I have not been able to find the cause of this problem, can you please help me analyze it?
class MMDetection(LabelStudioMLBase):
def __init__(self, config_file='D:/Code/AI/mmdetection/configs/faster_rcnn/faster-rcnn_r50_fpn_1x_coco.py',
checkpoint_file='D:/Code/AI/mmdetection/demo/checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth',
image_dir='D:/Code/AI/mmdetection/demo/laeblstudioml',
labels_file='{"airplane":"airplane","car":"car","person":"person"}', score_threshold=0.1, device='cpu', **kwargs):
super(MMDetection, self).__init__(**kwargs)
# config_file = config_file or os.environ['config_file']
# checkpoint_file = checkpoint_file or os.environ['checkpoint_file']
self.config_file = config_file
self.checkpoint_file = checkpoint_file
self.labels_file = labels_file
self.endpoint_url = kwargs.get('endpoint_url')
if self.endpoint_url:
logger.info(f'Using s3 endpoint url {self.endpoint_url}')
# default Label Studio image upload folder
upload_dir = os.path.join(get_data_dir(), 'media', 'upload')
self.image_dir = image_dir or upload_dir
logger.debug(f'{self.__class__.__name__} reads images from {self.image_dir}')
if self.labels_file and os.path.exists(self.labels_file):
self.label_map = json_load(self.labels_file)
else:
self.label_map = {}
self.from_name, self.to_name, self.value, self.labels_in_config = get_single_tag_keys(
self.parsed_label_config, 'RectangleLabels', 'Image')
schema = list(self.parsed_label_config.values())[0]
self.labels_in_config = set(self.labels_in_config)
# Collect label maps from `predicted_values="airplane,car"` attribute in <Label> tag
self.labels_attrs = schema.get('labels_attrs')
if self.labels_attrs:
for label_name, label_attrs in self.labels_attrs.items():
for predicted_value in label_attrs.get('predicted_values', '').split(','):
self.label_map[predicted_value] = label_name
print('Load new model from: ', config_file, checkpoint_file)
self.model = init_detector(config_file, checkpoint_file, device=device)
self.score_thresh = score_threshold
# print(f'debug {self.model}')
logger.debug(f'{self.model}')
def _get_image_url(self, task):
image_url = task['data'].get(self.value) or task['data'].get(DATA_UNDEFINED_NAME)
if image_url.startswith('s3://'):
# presign s3 url
r = urlparse(image_url, allow_fragments=False)
bucket_name = r.netloc
key = r.path.lstrip('/')
client = boto3.client('s3', endpoint_url=self.endpoint_url)
try:
image_url = client.generate_presigned_url(
ClientMethod='get_object',
Params={'Bucket': bucket_name, 'Key': key}
)
except ClientError as exc:
logger.warning(f'Can\'t generate presigned URL for {image_url}. Reason: {exc}')
return image_url
def predict(self, tasks, **kwargs):
# assert len(tasks) == 1
task = tasks[0]
image_url = self._get_image_url(task)
image_path = self.get_local_path(image_url)
model_results = inference_detector(self.model, image_path)
results = []
all_scores = []
img_width, img_height = get_image_size(image_path)
for bboxes, label in zip(model_results, self.model.CLASSES):
output_label = self.label_map.get(label, label)
if output_label not in self.labels_in_config:
print(output_label + ' label not found in project config.')
continue
for bbox in bboxes:
bbox = list(bbox)
if not bbox:
continue
score = float(bbox[-1])
if score < self.score_thresh:
continue
x, y, xmax, ymax = bbox[:4]
results.append({
'from_name': self.from_name,
'to_name': self.to_name,
'type': 'rectanglelabels',
'value': {
'rectanglelabels': [output_label],
'x': x / img_width * 100,
'y': y / img_height * 100,
'width': (xmax - x) / img_width * 100,
'height': (ymax - y) / img_height * 100
},
'score': score
})
all_scores.append(score)
avg_score = sum(all_scores) / max(len(all_scores), 1)
return [{
'result': results,
'score': avg_score
}]
def json_load(file, int_keys=False):
with io.open(file, encoding='utf8') as f:
data = json.load(f)
if int_keys:
return {int(k): v for k, v in data.items()}
else:
return data