(model_id: str, image_file_name: str)
| 235 | return str(result) |
| 236 | @task |
| 237 | def object_detection(model_id: str, image_file_name: str) -> str: |
| 238 | inference = InferenceApi(repo_id=model_id, token=CONFIG["huggingface"]["token"]) |
| 239 | img_data = image_to_bytes(f"{DIRPATH}/{INPUT_PATH}/{image_file_name}") |
| 240 | predicted = inference(data=img_data) |
| 241 | image = Image.open(BytesIO(img_data)) |
| 242 | draw = ImageDraw.Draw(image) |
| 243 | labels = list(item['label'] for item in predicted) |
| 244 | color_map = {} |
| 245 | for label in labels: |
| 246 | if label not in color_map: |
| 247 | color_map[label] = (random.randint(0, 255), random.randint(0, 100), random.randint(0, 255)) |
| 248 | for label in predicted: |
| 249 | box = label["box"] |
| 250 | draw.rectangle(((box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])), outline=color_map[label["label"]], width=2) |
| 251 | draw.text((box["xmin"]+5, box["ymin"]-15), label["label"], fill=color_map[label["label"]]) |
| 252 | name = str(uuid.uuid4())[:4] |
| 253 | image.save(f"{DIRPATH}/{OUTPUT_PATH}/{name}.jpg") |
| 254 | result = {} |
| 255 | result["generated image"] = f"{name}.jpg" |
| 256 | result["predicted"] = predicted |
| 257 | return str(result) |
| 258 | @task |
| 259 | def image_classification(model_id: str, image_file_name: str) -> str: |
| 260 | inference = InferenceApi(repo_id=model_id, token=CONFIG["huggingface"]["token"]) |
nothing calls this directly
no test coverage detected