Spaces:
Runtime error
Runtime error
File size: 3,425 Bytes
7ba936e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import cv2
import uuid
import os
from model import U2NET
from torch.autograd import Variable
from skimage import io, transform
from PIL import Image
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
app = FastAPI()
# Get The Current Directory
currentDir = os.path.dirname(__file__)
# Functions:
# Save Results
def save_output(image_name, output_name, pred, d_dir, type):
predict = pred
predict = predict.squeeze()
predict_np = predict.cpu().data.numpy()
im = Image.fromarray(predict_np*255).convert('RGB')
image = io.imread(image_name)
imo = im.resize((image.shape[1], image.shape[0]))
pb_np = np.array(imo)
if type == 'image':
# Make and apply mask
mask = pb_np[:, :, 0]
mask = np.expand_dims(mask, axis=2)
imo = np.concatenate((image, mask), axis=2)
imo = Image.fromarray(imo, 'RGBA')
imo.save(d_dir+output_name)
# Remove Background From Image (Generate Mask, and Final Results)
@app.get("/removeBG/")
async def removeBG(image_file: str):
# ------- Load Trained Model --------
print("---Loading Model---")
model_name = 'u2net'
model_dir = os.path.join(currentDir, 'saved_models',
model_name, model_name + '.pth')
net = U2NET(3, 1)
if torch.cuda.is_available():
net.load_state_dict(torch.load(model_dir))
net.cuda()
else:
net.load_state_dict(torch.load(model_dir, map_location='cpu'))
# ------- Load Trained Model --------
inputs_dir = os.path.join(currentDir, 'static/inputs/')
results_dir = os.path.join(currentDir, 'static/results/')
masks_dir = os.path.join(currentDir, 'static/masks/')
# convert string of image data to uint8
with open(image_file, "rb") as image:
f = image.read()
img = bytearray(f)
nparr = np.frombuffer(img, np.uint8)
if len(nparr) == 0:
return '---Empty image---'
# decode image
try:
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
except:
# build a response dict to send back to client
return "---Empty image---"
# save image to inputs
unique_filename = str(uuid.uuid4())
cv2.imwrite(inputs_dir+unique_filename+'.jpg', img)
# processing
image = transform.resize(img, (320, 320), mode='constant')
tmpImg = np.zeros((image.shape[0], image.shape[1], 3))
tmpImg[:, :, 0] = (image[:, :, 0]-0.485)/0.229
tmpImg[:, :, 1] = (image[:, :, 1]-0.456)/0.224
tmpImg[:, :, 2] = (image[:, :, 2]-0.406)/0.225
tmpImg = tmpImg.transpose((2, 0, 1))
tmpImg = np.expand_dims(tmpImg, 0)
image = torch.from_numpy(tmpImg)
image = image.type(torch.FloatTensor)
image = Variable(image)
d1, d2, d3, d4, d5, d6, d7 = net(image)
pred = d1[:, 0, :, :]
ma = torch.max(pred)
mi = torch.min(pred)
dn = (pred-mi)/(ma-mi)
pred = dn
save_output(inputs_dir+unique_filename+'.jpg', unique_filename +
'.png', pred, results_dir, 'image')
save_output(inputs_dir+unique_filename+'.jpg', unique_filename +
'.png', pred, masks_dir, 'mask')
return "---Success---"
#print("---Removing Background...")
# ------- Call The removeBg Function --------
# imgPath = "1.jpg" # Change this to your image path
# print(removeBg(imgPath))
|