Spaces:
Sleeping
Sleeping
File size: 1,920 Bytes
6307f85 26c4598 6307f85 26c4598 6307f85 26c4598 6307f85 26c4598 6307f85 26c4598 |
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 |
import os
import logging
from transformers import AutoImageProcessor, AutoModelForObjectDetection
from label_studio_ml.model import LabelStudioMLBase
from lxml import etree
class Model(LabelStudioMLBase):
image_processor = AutoImageProcessor.from_pretrained("diegokauer/conditional-detr-coe-int")
model = AutoModelForObjectDetection.from_pretrained("diegokauer/conditional-detr-coe-int")
def __init__(self, **kwargs):
# don't forget to call base class constructor
super(Model, self).__init__(**kwargs)
# you can preinitialize variables with keys needed to extract info from tasks and annotations and form predictions
self.model = model
self.tokenizer = image_processor
self.id2label = model.config.id2label
def predict(self, tasks, **kwargs):
""" This is where inference happens: model returns
the list of predictions based on input list of tasks
"""
predictions = []
for task in tasks:
predictions.append({
'score': 0.987, # prediction overall score, visible in the data manager columns
'model_version': 'delorean-20151021', # all predictions will be differentiated by model version
'result': [{
'from_name': self.from_name,
'to_name': self.to_name,
'type': 'choices',
'score': 0.5, # per-region score, visible in the editor
'value': {
'choices': [self.labels[0]]
}
}]
})
return predictions
def fit(self, annotations, **kwargs):
""" This is where training happens: train your model given list of annotations,
then returns dict with created links and resources
"""
return {'path/to/created/model': 'my/model.bin'} |