feat sketch 提取接口

fix
This commit is contained in:
zhouchengrong
2024-08-15 10:24:19 +08:00
parent 8b876b103f
commit 7ff3a72d8c
3 changed files with 5 additions and 4 deletions

View File

@@ -1,5 +1,6 @@
import json
import logging
import time
from fastapi import APIRouter, HTTPException
@@ -27,9 +28,11 @@ def image2sketch(request_item: Image2SketchModel):
}
"""
try:
start_time = time.time()
logger.info(f"image2sketch request item is : @@@@@@:{json.dumps(request_item.dict())}")
service = Image2SketchServer(request_item)
sketch_url = service.get_result()
logger.info(f"run time is : {time.time() - start_time}")
except Exception as e:
logger.warning(f"image2sketch Run Exception @@@@@@:{e}")
raise HTTPException(status_code=404, detail=str(e))

View File

@@ -7,7 +7,6 @@ class Config:
self.dataroot = "app/service/image2sketch/datasets/ref_unpair"
self.name = 'semi_unpair'
self.gpu_ids = [0]
self.checkpoints_dir = 'app/service/image2sketch/checkpoints/'
# 模型参数
self.model = 'unpaired'
self.input_nc = 3
@@ -48,5 +47,7 @@ class Config:
self.morm = 'batch'
if DEBUG:
self.style_image = "service/image2sketch/datasets/ref_unpair/testC/20180422151845_stEe4.jpeg"
self.checkpoints_dir = 'service/image2sketch/checkpoints/'
else:
self.checkpoints_dir = 'app/service/image2sketch/checkpoints/'
self.style_image = "app/service/image2sketch/datasets/ref_unpair/testC/20180422151845_stEe4.jpeg"

View File

@@ -55,9 +55,6 @@ class Image2SketchServer:
self.data['A'] = transform(A)
self.data['A'] = self.data['A'].unsqueeze(0).to(device)
def __del__(self):
torch.cuda.empty_cache()
def get_result(self):
self.model.set_input(self.data)
self.model.test() # run inference