jhynsoo/worlabel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Feat: WebSocket 연결

Browse Source
This commit is contained in:
김용수 2024-09-12 17:03:10 +09:00
parent c7d613365b
commit 6c46b64ce2
12 changed files with 73 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
backend/src/main/java/com/worlabel/domain/image/entity/Image.java
View File

@ -9,9 +9,6 @@ import lombok.AccessLevel;
import lombok.Getter;
import lombok.NoArgsConstructor;
import java.util.ArrayList;
import java.util.List;
@Getter
@Entity
@Table(name = "project_image")
@ -49,7 +46,7 @@ public class Image extends BaseEntity {
*/
@Column(name = "status", nullable = false)
@Enumerated(EnumType.STRING)
private LabelStatus status = LabelStatus.PENDING;
private LabelStatus status = LabelStatus.Pending;
/**
* 속한 폴더
@ -62,7 +59,7 @@ public class Image extends BaseEntity {
/**
* 이미지에 연결된 레이블
*/
@OneToOne(mappedBy = "image", cascade = CascadeType.ALL, orphanRemoval = true)
@OneToOne(mappedBy = "image", fetch = FetchType.LAZY, cascade = CascadeType.ALL, orphanRemoval = true)
private Label label;
private Image(final String imageTitle, final String imageUrl, final Integer order, final Folder folder) {

2
backend/src/main/java/com/worlabel/domain/image/entity/LabelStatus.java
View File

@ -4,7 +4,7 @@ import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonValue;
public enum LabelStatus {
PENDING,
Pending,
IN_PROGRESS,
NEED_REVIEW,
COMPLETED;

6
backend/src/main/java/com/worlabel/domain/image/repository/ImageRepository.java
View File

@ -12,10 +12,8 @@ public interface ImageRepository extends JpaRepository<Image, Long> {
Optional<Image> findByIdAndFolderId(Long imageId, Integer folderId);
// TODO: N + 1
@Query("select i from Image i " +
"join fetch i.folder f " +
"join fetch f.project p " +
"where p.id = :projectId")
"join fetch i.label l " +
"where i.folder.project.id = :projectId")
List<Image> findImagesByProjectId(@Param("projectId") Integer projectId);
}

4
backend/src/main/java/com/worlabel/domain/label/controller/LabelController.java
View File

@ -26,7 +26,6 @@ public class LabelController {
private final LabelService labelService;
@Operation(summary = "프로젝트 단위 오토레이블링", description = "해당 프로젝트 이미지를 오토레이블링합니다.")
@SwaggerApiSuccess(description = "해당 프로젝트가 오토 레이블링 됩니다.")
@SwaggerApiError({ErrorCode.EMPTY_REQUEST_PARAMETER, ErrorCode.SERVER_ERROR})
@ -51,7 +50,4 @@ public class LabelController {
labelService.save(imageId);
return SuccessResponse.empty();
}
}

11
backend/src/main/java/com/worlabel/domain/label/entity/Label.java
View File

@ -34,14 +34,6 @@ public class Label extends BaseEntity {
@JoinColumn(name = "image_id")
private Image image;
/**
* 속한 카테고리
* TODO: 한 레이블 카테고리에 속한걸 찾는데에 Json파일에 담기 때문에 카테고리는 Label Entity에 없어도 될 것 같음
*/
// @ManyToOne(fetch = FetchType.LAZY)
// @JoinColumn(name = "label_category_id")
// private LabelCategory labelCategory;
public static Label of(String jsonUrl, Image image) {
Label label = new Label();
label.url = jsonUrl;
@ -49,4 +41,7 @@ public class Label extends BaseEntity {
return label;
}
public void changeUrl(String newUrl){
}
}

2
backend/src/main/java/com/worlabel/domain/label/repository/LabelRepository.java
View File

@ -8,4 +8,6 @@ import java.util.Optional;
public interface LabelRepository extends JpaRepository<Label, Long> {
Optional<Label> findByImageId(Long imageId);
boolean existsByImageId(Long imageId);
}

28
backend/src/main/java/com/worlabel/domain/label/service/LabelService.java
View File

@ -59,20 +59,21 @@ public class LabelService {
List<Image> imageList = imageRepository.findImagesByProjectId(projectId);
List<ImageRequest> imageRequestList = imageList.stream().map(ImageRequest::of).toList();
AutoLabelingRequest autoLabelingRequest = AutoLabelingRequest.of(projectId, imageRequestList);
sendRequestToApi(autoLabelingRequest, projectType.getValue(), projectId);
}
List<AutoLabelingResponse> autoLabelingResponseList = sendRequestToApi(autoLabelingRequest, projectType.getValue(), projectId);
for (int index = 0; index < autoLabelingResponseList.size(); index++) {
AutoLabelingResponse response = autoLabelingResponseList.get(index);
Image image = imageList.get(index);
String uploadedUrl = s3UploadService.uploadJson(response.getData(), response.getImageUrl());
public void saveLabel(final AutoLabelingResponse autoLabelingResponse) {
String uploadUrl = s3UploadService.uploadJson(autoLabelingResponse.getData(), autoLabelingResponse.getImageUrl());
Image image = imageRepository.findById(autoLabelingResponse.getImageId())
.orElseThrow(() -> new CustomException(ErrorCode.IMAGE_NOT_FOUND));
Label label = labelRepository.findByImageId(response.getImageId())
.orElseGet(() -> Label.of(uploadedUrl, image));
if (!labelRepository.existsByImageId(autoLabelingResponse.getImageId())) {
Label label = Label.of(uploadUrl, image);
labelRepository.save(label);
}
}
private List<AutoLabelingResponse> sendRequestToApi(AutoLabelingRequest autoLabelingRequest, String apiEndpoint, int projectId) {
private void sendRequestToApi(AutoLabelingRequest autoLabelingRequest, String apiEndpoint, int projectId) {
String url = createApiUrl(apiEndpoint);
// RestTemplate을 동적으로 생성하여 사용
@ -94,9 +95,7 @@ public class LabelService {
String responseBody = Optional.ofNullable(response.getBody())
.orElseThrow(() -> new CustomException(ErrorCode.AI_SERVER_ERROR));
log.info("AI 서버 응답 -> {}", response.getBody());
return parseAutoLabelingResponseList(responseBody);
// return parseAutoLabelingResponseList(responseBody);
} catch (Exception e) {
log.error("AI 서버 요청 중 오류 발생: ", e);
throw new CustomException(ErrorCode.AI_SERVER_ERROR);
@ -107,17 +106,16 @@ public class LabelService {
JsonElement jsonElement = JsonParser.parseString(responseBody);
List<AutoLabelingResponse> autoLabelingResponseList = new ArrayList<>();
for (JsonElement element : jsonElement.getAsJsonArray()) {
AutoLabelingResponse response = parseAutoLabelingResponse(element);
AutoLabelingResponse response = parseAutoLabelingResponse(element.getAsJsonObject());
autoLabelingResponseList.add(response);
}
return autoLabelingResponseList;
}
/**
* jsonElement -> AutoLabelingResponse
* jsonObject -> AutoLabelingResponse
*/
private AutoLabelingResponse parseAutoLabelingResponse(JsonElement element) {
JsonObject jsonObject = element.getAsJsonObject();
public AutoLabelingResponse parseAutoLabelingResponse(JsonObject jsonObject) {
Long imageId = jsonObject.get("image_id").getAsLong();
String imageUrl = jsonObject.get("image_url").getAsString();
JsonObject data = jsonObject.get("data").getAsJsonObject();

2
backend/src/main/java/com/worlabel/domain/participant/repository/ParticipantRepository.java
View File

@ -26,6 +26,8 @@ public interface ParticipantRepository extends JpaRepository<Participant, Intege
@Param("projectId") Integer projectId);
Optional<Participant> findByMemberIdAndProjectId(Integer memberId, Integer projectId);
}

23
backend/src/main/java/com/worlabel/domain/progress/ProgressController.java
View File

@ -1,5 +1,11 @@
package com.worlabel.domain.progress;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.worlabel.domain.label.entity.dto.AutoLabelingResponse;
import com.worlabel.domain.label.service.LabelService;
import com.worlabel.global.service.S3UploadService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.SendTo;
@ -7,8 +13,13 @@ import org.springframework.stereotype.Controller;
@Slf4j
@Controller
@RequiredArgsConstructor
public class ProgressController {
private final LabelService labelService;
private final S3UploadService s3UploadService;
private final Gson gson;
@MessageMapping("/ai/train/progress")
@SendTo("/topic/progress")
public String handleTrainingProgress(String message) {
@ -16,4 +27,16 @@ public class ProgressController {
log.debug("Received message: {}", message);
return message;
}
@MessageMapping("/ai/predict/progress")
@SendTo("/topic/progress")
public String handlePredictProgress(String message) {
JsonObject jsonObject = gson.fromJson(message, JsonObject.class);
int progress = jsonObject.get("progress").getAsInt();
AutoLabelingResponse autoLabelingResponse = labelService.parseAutoLabelingResponse(jsonObject.getAsJsonObject("result"));
labelService.saveLabel(autoLabelingResponse);
log.debug("오토 레이블링 진행률 : {}%",progress);
return String.valueOf(progress);
}
}

2
backend/src/main/java/com/worlabel/domain/project/controller/ProjectController.java
View File

@ -83,8 +83,6 @@ public class ProjectController {
public BaseResponse<Void> trainModel(
@CurrentUser final Integer memberId,
@PathVariable("project_id") final Integer projectId) {
log.debug("훈련 요청 ");
projectService.train(memberId, projectId);
return SuccessResponse.empty();
}

41
backend/src/main/java/com/worlabel/domain/project/service/ProjectService.java
View File

@ -8,7 +8,6 @@ import com.worlabel.domain.participant.entity.WorkspaceParticipant;
import com.worlabel.domain.participant.entity.dto.ParticipantRequest;
import com.worlabel.domain.participant.repository.ParticipantRepository;
import com.worlabel.domain.participant.repository.WorkspaceParticipantRepository;
import com.worlabel.domain.project.dto.RequestDto;
import com.worlabel.domain.project.dto.RequestDto.TrainRequest;
import com.worlabel.domain.project.entity.Project;
import com.worlabel.domain.project.entity.dto.ProjectRequest;
@ -18,26 +17,16 @@ import com.worlabel.domain.workspace.entity.Workspace;
import com.worlabel.domain.workspace.repository.WorkspaceRepository;
import com.worlabel.global.exception.CustomException;
import com.worlabel.global.exception.ErrorCode;
import com.worlabel.global.service.AIWebSocketClient;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.ResponseEntity;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHttpHeaders;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.util.UriComponentsBuilder;
import java.net.URI;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
@Slf4j
@Service
@ -52,6 +41,12 @@ public class ProjectService {
private final WorkspaceParticipantRepository workspaceParticipantRepository;
private final RestTemplate restTemplate;
/**
* AI SERVER 주소
*/
@Value("${ai.server}")
private String aiServer;
public ProjectResponse createProject(final Integer memberId, final Integer workspaceId, final ProjectRequest projectRequest) {
Workspace workspace = getWorkspace(memberId, workspaceId);
Member member = getMember(memberId);
@ -128,14 +123,17 @@ public class ProjectService {
participantRepository.delete(participant);
}
@Async
public CompletableFuture<Void> train(final Integer memberId, final Integer projectId) {
public void train(final Integer memberId, final Integer projectId) {
// 멤버 권한 체크
checkEditorParticipant(memberId, projectId);
// 레디스 train 테이블에 존재하는지 확인
// TODO: 레디스 train 테이블에 존재하는지 확인 -> 이미 있으면 있다고 예외를 던져준다.
/*
없으면 redis 상태 테이블을 만든다. progressTable
*/
// FastAPI 서버로 학습 요청을 전송
String url = "http://localhost:8000/api/detection/train";
String url = aiServer + "/detection/train";
TrainRequest trainRequest = new TrainRequest();
trainRequest.setProjectId(projectId);
@ -145,12 +143,11 @@ public class ProjectService {
try {
ResponseEntity<String> result = restTemplate.postForEntity(url, trainRequest, String.class);
log.debug("응답 결과 {} ",result);
System.out.println("FastAPI 서버에 학습 요청을 성공적으로 전송했습니다. Project ID: " + projectId);
log.debug("FastAPI 서버에 학습 요청을 성공적으로 전송했습니다. Project ID: {}", projectId);
} catch (Exception e) {
System.err.println("FastAPI 서버에 학습 요청을 전송하는 중 오류가 발생했습니다: " + e.getMessage());
log.error("FastAPI 서버에 학습 요청을 전송하는 중 오류가 발생했습니다 {}",e.getMessage());
throw new CustomException(ErrorCode.AI_SERVER_ERROR);
}
return CompletableFuture.completedFuture(null);
}
private Workspace getWorkspace(final Integer memberId, final Integer workspaceId) {
@ -168,6 +165,12 @@ public class ProjectService {
.orElseThrow(() -> new CustomException(ErrorCode.PROJECT_NOT_FOUND));
}
private void checkEditorParticipant(final Integer memberId, final Integer projectId) {
if(participantRepository.doesParticipantUnauthorizedExistByMemberIdAndProjectId(memberId,projectId)){
throw new CustomException(ErrorCode.PARTICIPANT_UNAUTHORIZED);
}
}
private void checkExistParticipant(final Integer memberId, final Integer projectId) {
if (!participantRepository.existsByMemberIdAndProjectId(memberId, projectId)) {
throw new CustomException(ErrorCode.PARTICIPANT_UNAUTHORIZED);

5
backend/src/main/java/com/worlabel/global/service/S3UploadService.java
View File

@ -50,8 +50,7 @@ public class S3UploadService {
// String targetUrl = projectId + "/" + title + ".json"; // S3에 업로드할 대상 URL
String targetUrl = removeExtension(getKeyFromImageAddress(imageUrl)) + ".json";
log.debug("주소 {}", targetUrl);
// log.debug("주소 {}", targetUrl);
try {
byte[] jsonBytes = json.getBytes(StandardCharsets.UTF_8);
ObjectMetadata metadata = new ObjectMetadata();
@ -65,7 +64,7 @@ public class S3UploadService {
}
URL uploadedUrl = amazonS3.getUrl(bucket, targetUrl);
log.debug("Uploaded JSON URL: {}", uploadedUrl);
// log.debug("Uploaded JSON URL: {}", uploadedUrl);
return uploadedUrl.toString(); // 업로드된 파일의 URL 반환
} catch (Exception e) {
log.error("JSON 업로드 중 오류 발생: ", e);