yolo / app.py
Kroy's picture
Update app.py
652aa0e
raw
history blame
2.58 kB
import warnings
warnings.filterwarnings("ignore")
## import necessary packages
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
# Create a new FastAPI app instance
app = FastAPI()
# Root directory of the project
ROOT_DIR = os.getcwd()
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
os.system("pip install pycocotools==2.0.0")
# Define a function to handle the GET request at `/generate`
# The generate() function is defined as a FastAPI route that takes a
# string parameter called text. The function generates text based on the # input using the pipeline() object, and returns a JSON response
# containing the generated text under the key "output"
@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>).
"""
# Use the pipeline to generate text from the given input text
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)
# Get results and save them
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() # im_bytes: image in binary for
# Return the generated text in a JSON response
return {"output": im_bytes}