|
import warnings |
|
warnings.filterwarnings("ignore") |
|
|
|
|
|
import os |
|
import io |
|
import sys |
|
import base64 |
|
import random |
|
import argparse |
|
import math |
|
import numpy as np |
|
|
|
from typing import Any, Union,Dict, List |
|
import numpy as np |
|
import requests |
|
from PIL import Image |
|
from imageio import imread |
|
from keras import backend as K |
|
|
|
import coco |
|
import utils |
|
import model as modellib |
|
import visualize |
|
from classes import class_names |
|
from fastapi import FastAPI |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
ROOT_DIR = os.getcwd() |
|
|
|
|
|
MODEL_DIR = os.path.join(ROOT_DIR, "logs") |
|
|
|
|
|
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5") |
|
os.system("pip install pycocotools==2.0.0") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@app.get("/generate") |
|
def generate(path: str): |
|
""" |
|
Using the text summarization pipeline from `transformers`, summerize text |
|
from the given input text. The model used is `philschmid/bart-large-cnn-samsum`, which |
|
can be found [here](<https://huggingface.co/philschmid/bart-large-cnn-samsum>). |
|
""" |
|
|
|
K.clear_session() |
|
|
|
if not os.path.exists(COCO_MODEL_PATH): |
|
utils.download_trained_weights(COCO_MODEL_PATH) |
|
|
|
class InferenceConfig(coco.CocoConfig): |
|
GPU_COUNT = 1 |
|
IMAGES_PER_GPU = 1 |
|
config = InferenceConfig() |
|
|
|
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config) |
|
model.load_weights(COCO_MODEL_PATH, by_name=True) |
|
r = requests.get(path, stream=True) |
|
img = Image.open(io.BytesIO(r.content)).convert('RGB') |
|
open_cv_image = np.array(img) |
|
image = open_cv_image |
|
|
|
results = model.detect([image], verbose=1) |
|
|
|
|
|
r = results[0] |
|
output_image = visualize.display_instances_and_save(image, |
|
r['rois'], r['masks'], r['class_ids'], class_names, r['scores']) |
|
|
|
|
|
image = Image.fromarray(output_image) |
|
im_file = io.BytesIO() |
|
image.save(im_file, format="JPEG") |
|
im_bytes = im_file.getvalue() |
|
|
|
return {"output": im_bytes} |
|
|
|
|
|
|