新增接口：图片分割

src/main/java/com/ai/da/common/utils/RedisUtil.java

@@ -490,6 +490,7 @@ public class RedisUtil {
         return maxLikeCount != null ? Integer.parseInt(maxLikeCount) : 0;
     }
 
+    public final static String IMAGE_SEGMENTATION = "ImageSegmentation:";
 
     public final static String STRIPE_EXCEPTION_LOG = "StripeException:";
 

src/main/java/com/ai/da/controller/ElementController.java

@@ -18,7 +18,10 @@ import org.springframework.web.multipart.MultipartFile;
 
 import javax.annotation.Resource;
 import javax.validation.Valid;
+import javax.validation.constraints.Pattern;
+import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Collectors;
 
 
 @Api(tags = "collection模块")
@@ -106,4 +109,25 @@ public class ElementController {
         return Response.success();
     }
 
+    @ApiOperation(value = "图片分割")
+    @PostMapping("/imageSegmentation")
+    public Response<List<CollectionElementVO>> selectedImageSeg(
+            @RequestPart(value = "files", required = false) MultipartFile[] files,
+            @RequestParam(value = "type", required = false) @Pattern(regexp = "sketch|product", message = "类型必须是sketch或product") String type,
+            @RequestParam(value = "id", required = false) Long id) {
+        // 过滤空文件
+        List<MultipartFile> nonEmptyFiles = Arrays.stream(files)
+                .filter(file -> !file.isEmpty())
+                .collect(Collectors.toList());
+        // 参数校验
+        if ((nonEmptyFiles.isEmpty()) && id == null) {
+            throw new BusinessException("必须提供文件上传或ID");
+        }
+        if (!nonEmptyFiles.isEmpty() && id != null) {
+            throw new BusinessException("不能同时提供文件上传和ID");
+        }
+
+        return Response.success(collectionElementService.selectedImageSeg(nonEmptyFiles, id, type));
+    }
+
 }

src/main/java/com/ai/da/mapper/primary/entity/Library.java

@@ -69,6 +69,8 @@ public class Library implements Serializable {
      */
     private Integer width;
 
+    private String segmentedData;
+
     /**
      * 创建时间
      */

src/main/java/com/ai/da/model/vo/CollectionElementVO.java

@@ -4,6 +4,8 @@ import io.swagger.annotations.ApiModel;
 import io.swagger.annotations.ApiModelProperty;
 import lombok.Data;
 
+import java.util.List;
+
 @Data
 @ApiModel("element-响应")
 public class CollectionElementVO {
@@ -42,4 +44,6 @@ public class CollectionElementVO {
 
     private String originalUrl;
 
+    private List<String> segmentedImages;
+
 }

src/main/java/com/ai/da/python/PythonService.java

@@ -21,10 +21,7 @@ import com.ai.da.python.vo.*;
 import com.ai.da.service.DesignHistoryService;
 import com.ai.da.service.PythonTAllInfoService;
 import com.ai.da.service.SysFileService;
-import com.alibaba.fastjson.JSON;
+import com.alibaba.fastjson.*;
-import com.alibaba.fastjson.JSONArray;
-import com.alibaba.fastjson.JSONException;
-import com.alibaba.fastjson.JSONObject;
 import com.alibaba.fastjson.serializer.SerializerFeature;
 import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
 import com.google.common.collect.Lists;
@@ -4188,4 +4185,58 @@ public class PythonService {
         //生成失败
         throw new BusinessException("brandDNAGenerate.interface.exception");
     }
+
+    public List<ImageSegmentation.ImageDate> imageSegmentation(ImageSegmentation imageSegmentation) {
+        OkHttpClient client = new OkHttpClient().newBuilder()
+                .connectTimeout(30, TimeUnit.SECONDS)
+                .pingInterval(5, TimeUnit.SECONDS)//websocket轮训间隔(单位:秒)
+                .readTimeout(60, TimeUnit.SECONDS)//读取超时(单位:秒)
+                .writeTimeout(60, TimeUnit.SECONDS)//写入超时(单位:秒)
+                .build();
+
+        MediaType mediaType = MediaType.parse("application/json");
+        RequestBody body = RequestBody.create(mediaType, JSON.toJSONString(imageSegmentation));
+
+        log.info("modifyModelProportion 请求地址： {}，\n 参数：{}", accessPythonIp + ":" + accessPythonPort + "/api/clothing_seg", JSON.toJSONString(imageSegmentation));
+        Request request = new Request.Builder()
+                .url(accessPythonIp + ":" + accessPythonPort + "/api/clothing_seg")
+                .method("POST", body)
+                .addHeader("Content-Type", "application/json")
+                .build();
+        Response response = null;
+        try {
+            response = client.newCall(request).execute();
+        } catch (IOException ioException) {
+            log.error("PythonService##imageSegmentation异常###{}", ExceptionUtil.getThrowableList(ioException));
+            throw new BusinessException("generate.interface.error");
+        }
+        int responseCode = response.code();
+        String bodyString;
+        try {
+            bodyString = response.body().string();
+            if (responseCode != HttpURLConnection.HTTP_OK) {
+                // 基本不会有除200以外的code
+                log.info("imageSegmentation 失败。 Response code {}", responseCode);
+                throw new BusinessException("imageSegmentation 失败。 Response code " + responseCode);
+            }
+            JSONObject jsonObject = JSON.parseObject(bodyString);
+            if (response.isSuccessful() && jsonObject.get("msg").equals("OK!")) {
+                log.info("imageSegmentation 结果 ： {}", jsonObject.get("data").toString());
+                List<ImageSegmentation.ImageDate> seg = JSONObject.parseObject(
+                        jsonObject.get("data").toString(),
+                        new TypeReference<List<ImageSegmentation.ImageDate>>() {}
+                );
+                return seg;
+            }else {
+                log.info("imageSegmentation 失败。 Response code {}", responseCode);
+                throw new BusinessException("imageSegmentation 失败。 Response code " + responseCode);
+            }
+        } catch (IOException e) {
+            log.error("imageSegmentation 失败; error message => {}", e.getMessage());
+            response.close();
+            throw new BusinessException("generate.interface.error");
+        } finally {
+            response.close();
+        }
+    }
 }

src/main/java/com/ai/da/python/vo/ImageSegmentation.java

@@ -0,0 +1,20 @@
+package com.ai.da.python.vo;
+
+import lombok.Data;
+
+import java.util.List;
+
+@Data
+public class ImageSegmentation {
+
+    private Long user_id;
+    private List<ImageDate> image_data;
+
+    @Data
+    public class ImageDate{
+        public String image_url;
+        public String image_type;
+        // 作为入参时一起传入会怎样？
+        public List<String> clothing_url;
+    }
+}

src/main/java/com/ai/da/service/CollectionElementService.java

@@ -5,6 +5,7 @@ import com.ai.da.mapper.primary.entity.LibraryModelPoint;
 import com.ai.da.model.dto.*;
 import com.ai.da.model.vo.*;
 import com.baomidou.mybatisplus.extension.service.IService;
+import org.springframework.web.multipart.MultipartFile;
 
 import java.util.List;
 
@@ -136,4 +137,6 @@ public interface CollectionElementService extends IService<CollectionElement> {
      */
     CollectionElement editLevel2Type(Long elementId, String level2Type, String designType);
 
+    List<CollectionElementVO> selectedImageSeg(List<MultipartFile> files, Long id, String type);
+
 }

src/main/java/com/ai/da/service/impl/CollectionElementServiceImpl.java

@@ -3,6 +3,7 @@ package com.ai.da.service.impl;
 import cn.hutool.core.collection.CollectionUtil;
 import com.ai.da.common.config.FileProperties;
 import com.ai.da.common.config.exception.BusinessException;
+import com.ai.da.common.constant.CommonConstant;
 import com.ai.da.common.context.UserContext;
 import com.ai.da.common.enums.*;
 import com.ai.da.common.response.ResultEnum;
@@ -17,13 +18,16 @@ import com.ai.da.model.enums.StyleEnum;
 import com.ai.da.model.vo.*;
 import com.ai.da.python.PythonService;
 import com.ai.da.python.vo.DesignPythonItem;
+import com.ai.da.python.vo.ImageSegmentation;
 import com.ai.da.service.*;
 import com.alibaba.fastjson.JSON;
+import com.alibaba.fastjson.JSONObject;
 import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
 import com.baomidou.mybatisplus.core.metadata.IPage;
 import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
 import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
 import com.google.common.collect.Lists;
+import com.google.gson.Gson;
 import io.netty.util.internal.StringUtil;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.BeanUtils;
@@ -33,6 +37,7 @@ import org.springframework.stereotype.Service;
 import org.springframework.transaction.annotation.Transactional;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
+import org.springframework.web.multipart.MultipartFile;
 
 import javax.annotation.Resource;
 import java.io.File;
@@ -78,6 +83,8 @@ public class CollectionElementServiceImpl extends ServiceImpl<CollectionElementM
     private String collectionElement;
     @Value("${minio.bucketName.gradient}")
     private String gradientBucketName;
+    @Value("${minio.bucketName.users}")
+    private String userBucketName;
     @Resource
     private RedisUtil redisUtil;
 
@@ -948,4 +955,95 @@ public class CollectionElementServiceImpl extends ServiceImpl<CollectionElementM
 
         return generateDetail;
     }
+
+    // 对于上传图片或者从library选择的图片进行图片分割
+    public List<CollectionElementVO> selectedImageSeg(List<MultipartFile> files, Long id, String type) {
+        Long accountId = UserContext.getUserHolder().getId();
+        List<CollectionElementVO> resp = new ArrayList<>();
+        List<ImageSegmentation.ImageDate> imageDates = new ArrayList<>();
+
+        boolean isUploadMode = !files.isEmpty();
+        Library library = null;
+
+        // 判断是否是上传的图片
+        if (isUploadMode) {
+            String objectName = accountId + "/ImageSegment/input";
+            for (MultipartFile file : files) {
+                String md5 = MD5Utils.encryptFile(file);
+                String segmentedResult = redisUtil.getFromString(RedisUtil.IMAGE_SEGMENTATION + md5);
+                // 判断上传的图片是否有已分割的数据
+                if (StringUtil.isNullOrEmpty(segmentedResult)) {
+                    String path = minioUtil.upload(userBucketName, objectName, file);
+                    ImageSegmentation.ImageDate imageDate = new ImageSegmentation().new ImageDate();
+                    imageDate.setImage_url(path);
+                    imageDate.setImage_type(type);
+                    imageDates.add(imageDate);
+                } else {
+                    ImageSegmentation.ImageDate imageData = JSONObject.parseObject(segmentedResult, ImageSegmentation.ImageDate.class);
+                    resp.add(createCollectionElementVO(accountId, null, null, imageData.getImage_url(), imageData.getClothing_url()));
+                }
+            }
+        } else if (Objects.nonNull(id)) {
+            library = libraryService.getById(id);
+            // 判断从library中选择的图片是否有分割数据
+            if (StringUtil.isNullOrEmpty(library.getSegmentedData())) {
+                ImageSegmentation.ImageDate imageDate = new ImageSegmentation().new ImageDate();
+                imageDate.setImage_url(library.getUrl());
+                imageDate.setImage_type(type);
+                imageDates.add(imageDate);
+            } else {
+                List<String> restoredList = Arrays.asList(library.getSegmentedData().split(","));
+                resp.add(createCollectionElementVO(accountId, id, library.getLevel1Type(), library.getUrl(), restoredList));
+            }
+        }
+
+        // 处理需要分割的图片
+        if (!imageDates.isEmpty()) {
+            // 准备图片分割的参数
+            ImageSegmentation imageSegmentation = new ImageSegmentation();
+            imageSegmentation.setUser_id(accountId);
+            imageSegmentation.setImage_data(imageDates);
+            // 图片分割
+            List<ImageSegmentation.ImageDate> segmented = pythonService.imageSegmentation(imageSegmentation);
+            // 处理图片分割结果
+            for (ImageSegmentation.ImageDate imageData : segmented) {
+                if (isUploadMode) {
+                    // 上传的图片需要添加到redis中存一天
+                    String key = RedisUtil.IMAGE_SEGMENTATION +
+                            MD5Utils.encryptFile
+                                    (minioUtil.getPreSignedUrl
+                                            (imageData.getImage_url(), CommonConstant.MINIO_IMAGE_EXPIRE_TIME), false);
+                    redisUtil.addToString(key, new Gson().toJson(imageData), CommonConstant.GENERATE_RESULT_EXPIRE_TIME);
+                    resp.add(createCollectionElementVO(accountId, null, null, imageData.getImage_url(), imageData.getClothing_url()));
+                } else {
+                    // 从library中选择的图片需要更新数据库中对应图片的分割数据
+                    String segmentedData = String.join(",", imageData.getClothing_url());
+                    library.setSegmentedData(segmentedData);
+                    library.setUpdateDate(new Date());
+                    libraryService.updateById(library);
+                    resp.add(createCollectionElementVO(accountId, id, library.getLevel1Type(), library.getUrl(), imageData.getClothing_url()));
+                }
+            }
+        }
+        return resp;
+    }
+
+    // 准备返回数据
+    private CollectionElementVO createCollectionElementVO(Long accountId, Long id, String level1Type,
+                                                          String imageUrl, List<String> clothingUrls) {
+        CollectionElementVO vo = new CollectionElementVO();
+        vo.setAccountId(accountId);
+        vo.setId(id);
+        vo.setLevel1Type(level1Type);
+        vo.setUrl(minioUtil.getPreSignedUrl(imageUrl, CommonConstant.MINIO_IMAGE_EXPIRE_TIME));
+        vo.setMinIOPath(imageUrl);
+
+        List<String> segUrls = new ArrayList<>();
+        for (String seg : clothingUrls) {
+            segUrls.add(minioUtil.getPreSignedUrl(seg, CommonConstant.MINIO_IMAGE_EXPIRE_TIME));
+        }
+        vo.setSegmentedImages(segUrls);
+        return vo;
+    }
+
 }