Update app.py

app.py

@@ -3,40 +3,22 @@ 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 io
+import base64
 import requests
+import tensorflow as tf
+import tensorflow_hub as hub
 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 helper import *
 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")
+hub_handle = 'https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2'
+hub_module = hub.load(hub_handle)
 
 
  
@@ -45,42 +27,39 @@ os.system("pip install pycocotools==2.0.0")
 # 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):
+def generate(features: Dict[str, 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
+    content_image_url = features['url']
+    if 'style_url' not in features:
+        style_image_url = 'https://upload.wikimedia.org/wikipedia/commons/c/c5/Edvard_Munch%2C_1893%2C_The_Scream%2C_oil%2C_tempera_and_pastel_on_cardboard%2C_91_x_73_cm%2C_National_Gallery_of_Norway.jpg'
+    else:
+        style_image_url = features['style_url']
 
-    results = model.detect([image], verbose=1)
+    # Load images
+    content_img_size = (500, 500)
+    style_img_size = (300, 300)
 
-    # 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'])
+    style_image = load_image(style_image_url, style_img_size)
+    content_image = load_image(content_image_url, content_img_size)
+    style_image = tf.nn.avg_pool(
+        style_image, ksize=[3, 3], strides=[1, 1], padding='SAME')
 
+    # Stylize content image with given style image.
+    outputs = hub_module(tf.constant(content_image),
+                              tf.constant(style_image))
+    stylized_image = outputs[0]
 
-    image = Image.fromarray(output_image)
+    # get PIL image and convert to base64
+    img = Image.fromarray(np.uint8(stylized_image.numpy()[0] * 255))
     im_file = io.BytesIO()
-    image.save(im_file, format="JPEG")
-    im_bytes = im_file.getvalue()  # im_bytes: image in binary for
+    img.save(im_file, format="PNG")
+    im_bytes = base64.b64encode(im_file.getvalue()).decode("utf-8")
     # Return the generated text in a JSON response
     return {"output": im_bytes}