Add files via upload

Author: mkang30

docker-compose.yml

@@ -13,12 +13,12 @@ services:
     build: streamlit/
     depends_on:
       - fastapi
-    ports: 
+    ports:
         - 8501:8501
     networks:
       - deploy_network
     container_name: streamlit
 
 networks:
   deploy_network:
-    driver: bridge
+    driver: bridge

fastapi/Dockerfile

@@ -0,0 +1,19 @@
+FROM tiangolo/uvicorn-gunicorn:python3.7
+
+RUN mkdir /fastapi
+
+COPY requirements.txt /fastapi
+
+ADD /detectron2 /detectron2
+
+WORKDIR /fastapi
+
+RUN pip install -r requirements.txt -f https://download.pytorch.org/whl/torch_stable.html
+
+RUN pip install -e ../detectron2
+
+COPY . /fastapi
+
+EXPOSE 8000
+
+CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]

fastapi/detectron2.zip

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Sat Jul  4 13:31:36 2020
+
+@author: minseongkang
+"""
+import numpy as np
+from detectron2 import model_zoo
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import MetadataCatalog
+from PIL import Image
+import urllib.request
+
+"""
+This class represents the Instance Segmentation model Mask R-CNN
+"""
+class Segmentor:
+    
+    def __init__(self):
+        #create a predictor
+        self._cfg = get_cfg()
+        self._predictor = self._makePredictor()
+        self._class = MetadataCatalog.get(self._cfg.DATASETS.TRAIN[0]).get("thing_classes")
+    
+    """
+    This method initalizes the model and configuration 
+    to return the predictor
+    """
+    def _makePredictor(self):
+        self._cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
+        self._cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
+        self._cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
+        return DefaultPredictor(self._cfg)
+    
+    """
+    This method takes an opencv image and perfroms instance segmentation
+    """
+    def predict(self, image):
+        return self._predictor(image)
+
+    """
+    This method takes an output of the model and returns the segmentation
+    map of the image
+    """
+    def segMap(self, image,output):
+        v = Visualizer(image[:, :, ::-1], MetadataCatalog.get(self._cfg.DATASETS.TRAIN[0]), scale=1.2)
+        v = v.draw_instance_predictions(output["instances"].to("cpu"))
+        r = Image.fromarray(v.get_image()[:,:,::-1])
+        return r
+    """
+    This method takes an output of the model and returns an array of images
+    of key objects in the input image.
+    """
+    def decompose(self, image, output):
+        r = dict()
+        count = dict()
+        for i in range(len(output["instances"].pred_boxes)):
+            box = output["instances"].pred_boxes[i].tensor.numpy()[0]
+            dim = (box[0],box[1],box[2],box[3])
+            mask = output["instances"].pred_masks[i]
+            image2 = image[:,:,::-1].copy()
+            for j in range(len(image)):
+                for k in range(len(image[j])):
+                    if mask[j][k]==False:
+                        for l in range(len(image2[j][k])):
+                            image2[j][k][l]=255
+            pic = Image.fromarray(image2)
+            pic = pic.crop(dim)
+            pic = self._transparent(pic)
+            cl = self._class[output["instances"].pred_classes[i]]
+            if cl in count:
+                count[cl]+=1
+            else:
+                count[cl]=1
+            r[cl+str(count[cl])]=pic
+        return r
+    """
+    Input: PIL image
+    Output: PIL image with transparent background
+    """
+    def _transparent(self,image):
+        r = image.convert("RGBA")
+        pixels = r.getdata()
+        newPixels = []
+        for i in pixels:
+            if i[0]==255 and i[1]==255 and i[2]==255:
+                newPixels.append((255,255,255,0))
+            else:
+                newPixels.append(i)
+        r.putdata(newPixels)
+        return r
+
+"""
+image = Image.open(urllib.request.urlopen("https://static.independent.co.uk/s3fs-public/thumbnails/image/2015/06/06/15/Chris-Pratt.jpg?w968h681"))
+image = np.array(image)
+image = image[:,:,::-1].copy()
+pred = Segmentor()
+output = pred.predict(image)
+
+final = pred.decompose(image,output)
+
+
+for pic in final:
+    print(pic)
+    final[pic].show()
+    input()
+"""
+
+
+       
+
+        
+

fastapi/maskrcnn.py

@@ -0,0 +1,9 @@
+fastapi==0.58.0
+pillow==7.1.2
+torch==1.5.0+cpu
+torchvision==0.6.0+cpu
+python-multipart==0.0.5
+aiofiles==0.4.0
+pycocotools>=2.0.1
+opencv-python==4.3.0.36
+

fastapi/requirements.txt

@@ -0,0 +1,33 @@
+from fastapi import FastAPI, File
+from starlette.responses import FileResponse
+from maskrcnn import Segmentor
+from zipfile import ZipFile
+from PIL import Image
+import io
+import numpy as np
+
+
+#model = get_segmentator()
+model = Segmentor()
+
+app = FastAPI(title="Image Decomposer",
+              description='''Get the image decomposed into instances using the neural
+              network model Mask R-CNN implemented in detectron2 library.''',
+              version="0.1.0",
+              )
+
+
+@app.post("/segmentation")
+def get_segmentation_map(file: bytes = File(...)):
+    image = Image.open(io.BytesIO(file)).convert("RGB")
+    image = np.array(image)
+    image = image[:,:,::-1].copy()
+    output = model.predict(image)
+    instances = model.decompose(image, output)
+
+    with ZipFile('cont.zip', 'w') as zipObj:
+        for key in instances:
+            instances[key].save(key+".png")
+            zipObj.write(key+".png")
+
+    return FileResponse('cont.zip')

fastapi/server.py

@@ -0,0 +1,15 @@
+FROM tiangolo/uvicorn-gunicorn:python3.7
+
+RUN mkdir /streamlit
+
+COPY requirements.txt /streamlit
+
+WORKDIR /streamlit
+
+RUN pip install -r requirements.txt
+
+COPY . /streamlit
+
+EXPOSE 8501
+
+CMD ["streamlit", "run", "ui.py"]

streamlit/Dockerfile

@@ -0,0 +1,3 @@
+streamlit==0.61.0 
+requests==2.24.0
+requests-toolbelt==0.9.1

streamlit/requirements.txt

@@ -0,0 +1,49 @@
+import streamlit as st
+from requests_toolbelt.multipart.encoder import MultipartEncoder
+import requests
+from PIL import Image
+import io
+import zipfile
+import glob
+import os
+
+
+st.title('Image Decomposer')
+
+# fastapi endpoint
+url = 'http://fastapi:8000'
+endpoint = '/segmentation'
+
+st.write('''This application decomposes the image into multiple images of
+objects in the original image. Implemented using Mask R-CNN model of
+detectron2 library. The app can be used to get rid of the background, to
+separate specific objects from the rest of the image, or/and image analysis.''') # description and instructions
+
+image = st.file_uploader('insert image')  # image upload widget
+
+
+def process(image, server_url: str):
+
+    m = MultipartEncoder(
+        fields={'file': ('filename', image, 'image/jpeg')}
+        )
+
+    r = requests.post(server_url,
+                      data=m,
+                      headers={'Content-Type': m.content_type},
+                      timeout=8000)
+
+    return r
+
+
+if st.button('Get segmentation map'):
+    r = process(image, url+endpoint)
+    z = zipfile.ZipFile(io.BytesIO(r.content))
+    z.extractall()
+    image_list = []
+    for filename in glob.glob('*.png'):
+        im=Image.open(filename)
+        image_list.append(im)
+    st.image(image_list, width=300) # output dyptich
+    for filename in glob.glob('*.png'):
+        os.remove(filename)